11import logging
2+ import os
3+ from types import TracebackType
24from typing import Any , Literal
35
6+ import click
47from rich .segment import Segment
8+ from rich .style import Style
9+ from rich .text import Text
510from rich_toolkit import RichToolkit , RichToolkitTheme
6- from rich_toolkit .styles import MinimalStyle , TaggedStyle
11+ from rich_toolkit .styles import BaseStyle , MinimalStyle , TaggedStyle
12+
13+ from fastapi_cloud_cli .utils .version_check import (
14+ DISABLE_VERSION_CHECK_ENV ,
15+ BackgroundVersionCheck ,
16+ )
717
818logger = logging .getLogger (__name__ )
19+ VERSION_CHECK_CONTEXT_KEY = "fastapi_cloud_cli.version_check"
920
1021
1122class FastAPIStyle (TaggedStyle ):
@@ -20,9 +31,17 @@ def _get_tag_segments(
2031 animation_status : Literal ["started" , "stopped" , "error" ] | None = None ,
2132 ) -> tuple [list [Segment ], int ]:
2233 if not is_animated :
23- return super ()._get_tag_segments (
34+ tag_segments , left_padding = super ()._get_tag_segments (
2435 metadata , is_animated , done , animation_status = animation_status
2536 )
37+ tag_style = metadata .get ("tag_style" )
38+ if isinstance (tag_style , (str , Style )):
39+ style = self .console .get_style (tag_style )
40+ tag_segments = [
41+ Segment (segment .text , style = style ) for segment in tag_segments
42+ ]
43+
44+ return tag_segments , left_padding
2645
2746 emojis = [
2847 "🥚" ,
@@ -47,14 +66,96 @@ def _get_tag_segments(
4766 return [Segment (tag )], left_padding
4867
4968
50- def get_rich_toolkit (minimal : bool = False ) -> RichToolkit :
69+ class FastAPIRichToolkit (RichToolkit ):
70+ def __init__ (
71+ self ,
72+ style : BaseStyle | None = None ,
73+ theme : RichToolkitTheme | None = None ,
74+ handle_keyboard_interrupts : bool = True ,
75+ print_spacing : bool = True ,
76+ ) -> None :
77+ super ().__init__ (
78+ style = style ,
79+ theme = theme ,
80+ handle_keyboard_interrupts = handle_keyboard_interrupts ,
81+ )
82+ self ._print_spacing = print_spacing
83+ self ._version_check : BackgroundVersionCheck | None = None
84+ self ._print_update_on_exit = False
85+
86+ def __enter__ (self ) -> "FastAPIRichToolkit" :
87+ self ._version_check = self ._get_version_check ()
88+
89+ if self ._print_spacing :
90+ self .console .print ()
91+ return self
92+
93+ def __exit__ (
94+ self ,
95+ exc_type : type [BaseException ] | None ,
96+ exc_value : BaseException | None ,
97+ traceback : TracebackType | None ,
98+ ) -> bool | None :
99+ is_keyboard_interrupt = exc_type is KeyboardInterrupt
100+
101+ if is_keyboard_interrupt and self ._version_check is not None :
102+ self ._version_check .suppress ()
103+ elif self ._print_update_on_exit :
104+ self ._print_update_message ()
105+
106+ if self ._print_spacing and not is_keyboard_interrupt :
107+ self .console .print ()
108+
109+ if self .handle_keyboard_interrupts and is_keyboard_interrupt :
110+ return True
111+
112+ return None
113+
114+ def _get_version_check (self ) -> BackgroundVersionCheck | None :
115+ if os .environ .get (DISABLE_VERSION_CHECK_ENV ) == "1" :
116+ return None
117+
118+ context = click .get_current_context (silent = True )
119+ if context is None :
120+ version_check = BackgroundVersionCheck ()
121+ version_check .start ()
122+ self ._print_update_on_exit = True
123+ return version_check
124+
125+ stored_version_check = context .meta .get (VERSION_CHECK_CONTEXT_KEY )
126+ if isinstance (stored_version_check , BackgroundVersionCheck ):
127+ return stored_version_check
128+
129+ version_check = BackgroundVersionCheck ()
130+ version_check .start ()
131+ context .meta [VERSION_CHECK_CONTEXT_KEY ] = version_check
132+ context .call_on_close (self ._print_update_message )
133+
134+ return version_check
135+
136+ def _print_update_message (self ) -> None :
137+ if self ._version_check is None :
138+ return
139+
140+ message = self ._version_check .get_update_message ()
141+ if message :
142+ self .print (Text .from_markup (message ), tag = "update" , tag_style = "tag.update" )
143+
144+
145+ def get_rich_toolkit (
146+ minimal : bool = False ,
147+ * ,
148+ print_spacing : bool = True ,
149+ handle_keyboard_interrupts : bool = True ,
150+ ) -> RichToolkit :
51151 style = MinimalStyle () if minimal else FastAPIStyle (tag_width = 11 )
52152
53153 theme = RichToolkitTheme (
54154 style = style ,
55155 theme = {
56156 "tag.title" : "white on #009485" ,
57157 "tag" : "white on #007166" ,
158+ "tag.update" : "black on yellow" ,
58159 "placeholder" : "grey62" ,
59160 "text" : "white" ,
60161 "selected" : "#007166" ,
@@ -65,4 +166,8 @@ def get_rich_toolkit(minimal: bool = False) -> RichToolkit:
65166 },
66167 )
67168
68- return RichToolkit (theme = theme )
169+ return FastAPIRichToolkit (
170+ theme = theme ,
171+ handle_keyboard_interrupts = handle_keyboard_interrupts ,
172+ print_spacing = print_spacing ,
173+ )
0 commit comments