22import sys
33from typing import Any
44
5- from rich_toolkit import RichToolkit , RichToolkitTheme
6- from rich_toolkit .styles import TaggedStyle
5+ from rich ._loop import loop_first
6+ from rich .console import Console , ConsoleOptions , RenderableType , RenderResult
7+ from rich .segment import Segment
8+ from rich .text import Text
9+ from rich_toolkit import RichToolkit
10+ from rich_toolkit .element import Element
11+ from rich_toolkit .styles import BaseStyle
712from uvicorn .logging import DefaultFormatter
813
14+ logger = logging .getLogger (__name__ )
15+
16+
17+ def should_use_rich_logs () -> bool :
18+ """Return True when stdout is a TTY and rich logs should be used, False otherwise."""
19+ return sys .stdout .isatty ()
20+
21+
22+ class IndentedBlock :
23+ """Indent a renderable, hanging a prefix (e.g. an emoji bullet) on the
24+ first line and aligning wrapped/extra lines under the text."""
25+
26+ def __init__ (
27+ self ,
28+ renderable : RenderableType ,
29+ * ,
30+ first_prefix : Text ,
31+ prefix : Text ,
32+ ) -> None :
33+ self .renderable = renderable
34+ self .first_prefix = first_prefix
35+ self .prefix = prefix
36+
37+ # Text renders its `end` ("\n" by default), which would break lines
38+ self .first_prefix .end = ""
39+ self .prefix .end = ""
40+
41+ def __rich_console__ (
42+ self , console : Console , options : ConsoleOptions
43+ ) -> RenderResult :
44+ prefix_width = max (self .first_prefix .cell_len , self .prefix .cell_len )
45+ lines = console .render_lines (
46+ self .renderable ,
47+ options .update_width (options .max_width - prefix_width ),
48+ pad = False ,
49+ )
50+
51+ new_line = Segment .line ()
52+
53+ for first , line in loop_first (lines ):
54+ if any (segment .text .strip () for segment in line ):
55+ yield from console .render (
56+ self .first_prefix if first else self .prefix , options
57+ )
58+ yield from line
59+
60+ yield new_line
61+
62+
63+ class FastAPIStyle (BaseStyle ):
64+ """Uniform left indent with emoji bullets hanging to the left of the text.
65+
66+ ``emoji=`` renders as a bullet in a fixed column; ``bullet=False`` drops
67+ the bullet column and just indents. This is a small, purpose-built subset
68+ of the richer style in fastapi-cloud-cli.
69+ """
70+
71+ content_padding = 1
72+ emoji_column_width = 3
73+
74+ def render_element (
75+ self ,
76+ element : Any ,
77+ is_active : bool = False ,
78+ done : bool = False ,
79+ parent : Element | None = None ,
80+ ** kwargs : Any ,
81+ ) -> RenderableType :
82+ rendered = super ().render_element (
83+ element = element , is_active = is_active , done = done , parent = parent , ** kwargs
84+ )
85+
86+ emoji = kwargs .get ("emoji" , "" )
87+
88+ if not emoji and not kwargs .get ("bullet" , True ):
89+ indent = Text (" " * (self .content_padding + 1 ))
90+ return IndentedBlock (rendered , first_prefix = indent , prefix = indent )
91+
92+ return IndentedBlock (
93+ rendered ,
94+ first_prefix = self ._get_bullet_prefix (emoji ),
95+ prefix = Text (" " * (self .content_padding + self .emoji_column_width )),
96+ )
97+
98+ def _get_bullet_prefix (self , emoji : str ) -> Text :
99+ prefix = Text (" " * self .content_padding )
100+
101+ if emoji :
102+ prefix .append_text (Text .from_markup (emoji ))
103+
104+ prefix .pad_right (
105+ self .content_padding + self .emoji_column_width - prefix .cell_len
106+ )
107+
108+ return prefix
109+
110+
111+ LOG_LEVEL_COLORS = {
112+ "debug" : "blue" ,
113+ "info" : "cyan" ,
114+ "warning" : "yellow" ,
115+ "warn" : "yellow" ,
116+ "error" : "red" ,
117+ "critical" : "magenta" ,
118+ "fatal" : "magenta" ,
119+ }
120+
121+
122+ def _get_log_bullet (level : str ) -> str :
123+ """Colored bar rendered in the emoji bullet column, matching the log
124+ level, like the ``fastapi cloud logs`` output."""
125+ color = LOG_LEVEL_COLORS .get (level .lower (), "dim" )
126+
127+ return f"[{ color } { color }
128+
9129
10130class CustomFormatter (DefaultFormatter ):
11131 def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
@@ -14,18 +134,15 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
14134
15135 def formatMessage (self , record : logging .LogRecord ) -> str :
16136 message = record .getMessage ()
17- result = self .toolkit .print_as_string (message , tag = record .levelname )
137+ result = self .toolkit .print_as_string (
138+ message , emoji = _get_log_bullet (record .levelname )
139+ )
18140 # Prepend newline to fix alignment after ^C is printed by the terminal
19141 if message == "Shutting down" :
20142 result = "\n " + result
21143 return result
22144
23145
24- def should_use_rich_logs () -> bool :
25- """Return True when stdout is a TTY and rich logs should be used, False otherwise."""
26- return sys .stdout .isatty ()
27-
28-
29146def get_uvicorn_log_config () -> dict [str , Any ]:
30147 return {
31148 "version" : 1 ,
@@ -69,23 +186,5 @@ def get_uvicorn_log_config() -> dict[str, Any]:
69186 }
70187
71188
72- logger = logging .getLogger (__name__ )
73-
74-
75189def get_rich_toolkit () -> RichToolkit :
76- theme = RichToolkitTheme (
77- style = TaggedStyle (tag_width = 11 ),
78- theme = {
79- "tag.title" : "white on #009485" ,
80- "tag" : "white on #007166" ,
81- "placeholder" : "grey85" ,
82- "text" : "white" ,
83- "selected" : "#007166" ,
84- "result" : "grey85" ,
85- "progress" : "on #007166" ,
86- "error" : "red" ,
87- "log.info" : "black on blue" ,
88- },
89- )
90-
91- return RichToolkit (theme = theme )
190+ return RichToolkit (style = FastAPIStyle ())
0 commit comments