def test_refresh_screen():
layout = Layout()
layout.split_row(Layout(name="foo"), Layout(name="bar"))
console = Console(force_terminal=True, width=20, height=5)
with console.capture():
console.print(layout)
with console.screen():
with console.capture() as capture:
layout.refresh_screen(console, "foo")
result = capture.get()
print()
print(repr(result))
expected = "\x1b[1;1H\x1b[34m╭─\x1b[0m\x1b[34m \x1b[0m\x1b[32m'foo'\x1b[0m\x1b[34m─╮\x1b[0m\x1b[2;1H\x1b[34m│\x1b[0m \x1b[1;35mLayout\x1b[0m \x1b[34m│\x1b[0m\x1b[3;1H\x1b[34m│\x1b[0m \x1b[1m(\x1b[0m \x1b[34m│\x1b[0m\x1b[4;1H\x1b[34m│\x1b[0m \x1b[33mna\x1b[0m \x1b[34m│\x1b[0m\x1b[5;1H\x1b[34m╰────────╯\x1b[0m"
assert result == expected
def test_capture():
console = Console()
with console.capture() as capture:
with pytest.raises(CaptureError):
capture.get()
console.print("Hello")
assert capture.get() == "Hello\n"
def rich_print(string: str,
console: Console,
style: str = "",
end: str = "\n"):
with console.capture() as capture:
console.print(string, style=style, end=end)
return capture.get()
def _format_output(args, code):
if not args['color']:
return code
lexer = None
# try to find a lexer using the StackOverflow tags
# or the query arguments
for keyword in args['query'].split() + args['tags']:
try:
lexer = get_lexer_by_name(keyword).name
break
except ClassNotFound:
pass
# no lexer found above, use the guesser
if not lexer:
try:
lexer = guess_lexer(code).name
except ClassNotFound:
return code
syntax = Syntax(code,
lexer,
background_color="default",
line_numbers=False)
console = Console(record=True)
with console.capture() as capture:
console.print(syntax)
return capture.get()
def test_rich_pretty_angular() -> None:
console = Console()
with console.capture() as capture:
console.print(Bar("hello", bar=3))
result = capture.get()
expected = "<Bar 'hello' 'hello' bar=3 egg=1>\n"
assert result == expected
def test_rich_pretty() -> None:
console = Console()
with console.capture() as capture:
console.print(Foo("hello", bar=3))
result = capture.get()
expected = "Foo('hello', 'hello', bar=3, egg=1)\n"
assert result == expected
def test_render():
layout = Layout(name="root")
repr(layout)
layout.split_column(Layout(name="top"), Layout(name="bottom"))
top = layout["top"]
top.update(Panel("foo"))
print(type(top._renderable))
assert isinstance(top.renderable, Panel)
layout["bottom"].split_row(Layout(name="left"), Layout(name="right"))
assert layout["root"].name == "root"
assert layout["left"].name == "left"
with pytest.raises(KeyError):
top["asdasd"]
layout["left"].update("foobar")
print(layout["left"].children)
console = Console(width=60, color_system=None)
with console.capture() as capture:
console.print(layout, height=10)
result = capture.get()
print(repr(result))
expected = "╭──────────────────────────────────────────────────────────╮\n│ foo │\n│ │\n│ │\n╰──────────────────────────────────────────────────────────╯\nfoobar ╭───── 'right' (30 x 5) ─────╮\n │ │\n │ Layout(name='right') │\n │ │\n ╰────────────────────────────╯\n"
assert result == expected
def print_to_file(self, log: Path):
with open(log, "w") as file:
tree = self._get_tree(verbose=True)
console = Console(record=True)
with console.capture() as capture:
console.print(tree)
file.write(console.export_text())
def on_task_exception(self):
logger.exception("An internal error occurred in the application")
toast_msg = ("应用发生内部错误" if "zh" in session_info.user_language else
"An internal error occurred in the application")
e_type, e_value, e_tb = sys.exc_info()
lines = traceback.format_exception(e_type, e_value, e_tb)
traceback_msg = "".join(lines)
traceback_console = Console(color_system="truecolor",
tab_size=2,
record=True,
width=90)
with traceback_console.capture(): # prevent logging to stdout again
traceback_console.print_exception(word_wrap=True,
extra_lines=1,
show_locals=True)
if State.theme == "dark":
theme = DARK_TERMINAL_THEME
else:
theme = LIGHT_TERMINAL_THEME
html = traceback_console.export_html(theme=theme,
code_format=TRACEBACK_CODE_FORMAT,
inline_styles=True)
try:
popup(title=toast_msg, content=put_html(html), size=PopupSize.LARGE)
run_js(
"console.error(traceback_msg)",
traceback_msg="Internal Server Error\n" + traceback_msg,
)
except Exception:
pass
def test_screen():
console = Console(color_system=None, width=20, height=5, legacy_windows=False)
with console.capture() as capture:
console.print(Screen("foo\nbar\nbaz\nfoo\nbar\nbaz\foo"))
result = capture.get()
print(repr(result))
expected = "foo \nbar \nbaz \nfoo \nbar "
assert result == expected
def test_tree():
layout = Layout(name="root")
layout.split(Layout("foo", size=2), Layout("bar"))
console = Console(width=60, color_system=None)
with console.capture() as capture:
console.print(layout.tree, height=10)
result = capture.get()
print(repr(result))
expected = "⬇ 'root' (ratio=1) \n├── ■ (size=2) \n└── ■ (ratio=1) \n"
assert result == expected
def test_tree():
layout = Layout(name="root")
layout.split(Layout("foo", size=2), Layout("bar", name="bar"))
layout["bar"].split_row(Layout(), Layout())
console = Console(width=60, color_system=None)
with console.capture() as capture:
console.print(layout.tree, height=10)
result = capture.get()
print(repr(result))
expected = "⬍ Layout(name='root') \n├── ⬍ Layout(size=2) \n└── ⬌ Layout(name='bar') \n ├── ⬍ Layout() \n └── ⬍ Layout() \n"
print(result, "\n", expected)
assert result == expected
def test_decode_example():
ansi_bytes = b"\x1b[01m\x1b[KC:\\Users\\stefa\\AppData\\Local\\Temp\\tmp3ydingba:\x1b[m\x1b[K In function '\x1b[01m\x1b[Kmain\x1b[m\x1b[K':\n\x1b[01m\x1b[KC:\\Users\\stefa\\AppData\\Local\\Temp\\tmp3ydingba:3:5:\x1b[m\x1b[K \x1b[01;35m\x1b[Kwarning: \x1b[m\x1b[Kunused variable '\x1b[01m\x1b[Ka\x1b[m\x1b[K' [\x1b[01;35m\x1b[K-Wunused-variable\x1b[m\x1b[K]\n 3 | int \x1b[01;35m\x1b[Ka\x1b[m\x1b[K=1;\n | \x1b[01;35m\x1b[K^\x1b[m\x1b[K\n"
ansi_text = ansi_bytes.decode("utf-8")
text = Text.from_ansi(ansi_text)
console = Console(force_terminal=True,
legacy_windows=False,
color_system="truecolor")
with console.capture() as capture:
console.print(text)
result = capture.get()
print(repr(result))
expected = "\x1b[1mC:\\Users\\stefa\\AppData\\Local\\Temp\\tmp3ydingba:\x1b[0m In function '\x1b[1mmain\x1b[0m':\n\x1b[1mC:\\Users\\stefa\\AppData\\Local\\Temp\\tmp3ydingba:3:5:\x1b[0m \x1b[1;35mwarning: \x1b[0munused variable '\x1b[1ma\x1b[0m' \n[\x1b[1;35m-Wunused-variable\x1b[0m]\n 3 | int \x1b[1;35ma\x1b[0m=1;\n | \x1b[1;35m^\x1b[0m\n"
assert result == expected
def test_capture_and_record(capsys):
recorder = Console(record=True)
recorder.print("ABC")
with recorder.capture() as capture:
recorder.print("Hello")
assert capture.get() == "Hello\n"
recorded_text = recorder.export_text()
out, err = capsys.readouterr()
assert recorded_text == "ABC\nHello\n"
assert capture.get() == "Hello\n"
assert out == "ABC\n"
class TqdmHandler(logging.Handler):
def __init__(self, level=logging.NOTSET) -> None:
super().__init__(level=level)
self.console = Console()
def emit(self, record: logging.LogRecord) -> None:
try:
msg = self.format(record)
# Colorize
with self.console.capture() as capture:
self.console.print(msg, end='')
msg = capture.get()
_tqdm.write(msg)
self.flush()
except Exception:
self.handleError(record)
def main(self, stdscr):
page = 0
while True:
stdscr.clear()
console = Console()
size = console.size.height - 5
# Pagination
items = self.paginate(page, size)
# Print on buffer
with console.capture() as capture:
console.print(ListTable(items))
output = capture.get()
from richer.text import AnsiEscapeText
tokens = AnsiEscapeText(output).tokens
for t in tokens:
mode = curses.A_NORMAL
for a in t.attr.split(';'):
mode = mode | self.graphic_mode.get(a, curses.A_NORMAL)
stdscr.addstr(t.row, t.col, t.text, mode)
stdscr.refresh()
while True:
c = stdscr.getch()
if c == ord('q'):
return
elif c == curses.KEY_LEFT or c == curses.KEY_PPAGE:
if page > 0:
page -= 1
break
elif c == curses.KEY_RIGHT or c == curses.KEY_NPAGE:
if page < self.max_page(size):
page += 1
break
def test_vertical_align_top():
console = Console(_environ={})
def make_table(vertical_align):
table = Table(show_header=False, box=box.SQUARE)
table.add_column(vertical=vertical_align)
table.add_row("foo", "\n".join(["bar"] * 5))
return table
with console.capture() as capture:
console.print(make_table("top"))
console.print()
console.print(make_table("middle"))
console.print()
console.print(make_table("bottom"))
console.print()
result = capture.get()
print(repr(result))
expected = "┌─────┬─────┐\n│ foo │ bar │\n│ │ bar │\n│ │ bar │\n│ │ bar │\n│ │ bar │\n└─────┴─────┘\n\n┌─────┬─────┐\n│ │ bar │\n│ │ bar │\n│ foo │ bar │\n│ │ bar │\n│ │ bar │\n└─────┴─────┘\n\n┌─────┬─────┐\n│ │ bar │\n│ │ bar │\n│ │ bar │\n│ │ bar │\n│ foo │ bar │\n└─────┴─────┘\n\n"
assert result == expected
def parse_md(md):
from rich.console import Console
from rich.markdown import Markdown
import re
table_re = r'(^|[^|]\n)((?:^\|[^\n]*\|(?:\n|$))+)([^|]|$)'
md = re.sub(r'\n([\r\t ]*\n)+', r'\n\n', md, flags=re.MULTILINE)
md = re.sub(r'^[\t ]*\|', r'|', md, flags=re.MULTILINE)
md = re.sub(r'\|[\t ]*$', r'|', md, flags=re.MULTILINE)
tables = []
md = re.sub(table_re,
lambda match: parse_md_table(match, tables),
md,
flags=re.MULTILINE | re.S)
console = Console(width=70)
with console.capture() as capture:
console.print(Markdown(md))
formated_text = capture.get()
for i in range(len(tables)):
formated_text = re.sub('<#MD-TABLE-' + str(i) + '#>', tables[i],
formated_text)
return re.sub(r'<br/?>', r'\n', formated_text).strip()
def return_formatted_ans(self, ques_id):
# This function uses uses Rich Markdown to format answers body.
body_markdown = self.answer_data[int(ques_id)]
body_markdown = str(body_markdown)
xml_markup_replacement = [
("&", "&"),
("<", "<"),
(">", ">"),
(""", '"'),
("'", "'"),
("'", "'"),
]
for convert_from, convert_to in xml_markup_replacement:
body_markdown = body_markdown.replace(convert_from, convert_to)
width = os.get_terminal_size().columns
console = Console(width=width - 4)
markdown = Markdown(body_markdown, hyperlinks=False)
with console.capture() as capture:
console.print(markdown)
highlighted = capture.get()
if not ("UTF" in locale.getlocale()[1]):
box_replacement = [
("─", "-"),
("═", "="),
("║", "|"),
("│", "|"),
("┌", "+"),
("└", "+"),
("┐", "+"),
("┘", "+"),
("╔", "+"),
("╚", "+"),
("╗", "+"),
("╝", "+"),
("•", "*"),
]
for convert_from, convert_to in box_replacement:
highlighted = highlighted.replace(convert_from, convert_to)
return highlighted
def parse_md_table(match, tables_memo):
from rich.console import Console
from rich.table import Table
from rich.style import Style
from rich.align import Align
from rich import box
import re
[table_header, table_body, columns] = split_md_table(match.group(2))
b = box.Box(' \n \n══╪═\n \n┈┈┼┈\n┈┈┼┈\n \n ')
table = Table(box=b,
show_header=False,
show_edge=False,
border_style=Style(color='#222222'))
for col_align in columns:
table.add_column(None, justify=col_align)
sty_header = Style(bgcolor='blue', color='#000000', bold=True)
sty_odd = Style(bgcolor="#111111", color='white')
sty_even = Style(bgcolor="#222222", color='white')
for row in table_header:
row = map(Align.center, row)
table.add_row(*row, style=sty_header)
num = 0
for row in table_body:
num += 1
style = sty_even if (num % 2 == 0) else sty_odd
table.add_row(*row, style=style)
console = Console(width=70)
with console.capture() as capture:
console.print(table)
formated_text = '\n'.join(capture.get().split('\n')[0:-1])
before = '' if match.group(1) == '\n\n' else '\n'
after = ' ' if match.group(3) == '\n' else ' \n'
tables_memo.append(before + formated_text)
return match.group(1) + \
'<#MD-TABLE-'+str(len(tables_memo)-1)+'#>' + after + \
match.group(3)
class ConsoleAndPanel:
"""Create a rich console to wrap the console print with a Panel"""
def __init__(self):
self.console = Console(theme=CUSTOM_THEME, highlight=False, soft_wrap=True)
def capture(self):
return self.console.capture()
@staticmethod
def filter_rich_tags(text):
for val in RICH_TAGS:
text = text.replace(val, "")
return text
def print(self, *args, **kwargs):
if kwargs and "text" in list(kwargs) and "menu" in list(kwargs):
if gtff.ENABLE_RICH:
if gtff.ENABLE_RICH_PANEL:
self.console.print(
panel.Panel(
kwargs["text"],
title=kwargs["menu"],
subtitle_align="right",
subtitle="Gamestonk Terminal",
)
)
else:
self.console.print(kwargs["text"])
else:
print(self.filter_rich_tags(kwargs["text"]))
else:
if gtff.ENABLE_RICH:
self.console.print(*args, **kwargs)
else:
print(*args, **kwargs)
def test_pprint(response):
console = Console(force_terminal=False, width=120)
with console.capture() as output:
pprint(response)
rendered = output.get()
def generate_preview(markdown: str):
console = Console(width=40)
with console.capture() as capture:
console.print(Markdown(markdown))
preview.text = ANSI(capture.get())
def generate_string_report(self):
recorder = Console(record=True)
with recorder.capture():
recorder.print(self.report)
return recorder.export_text()
def output(
self,
outdata: Union[List[str], Dict[str, Any]],
tablefmt: str = "rich",
title: str = None,
caption: str = None,
account: str = None,
config=None,
ok_status: Union[int, List[int], Tuple[int, str],
List[Tuple[int, str]]] = None,
) -> str:
# log.debugv(f"data passed to output():\n{pprint(outdata, indent=4)}")
def _do_subtables(data: list, tablefmt: str = "rich"):
out = []
for inner_dict in data: # the object: switch/vlan etc dict
for key, val in inner_dict.items():
if not isinstance(val, (list, dict, tuple)):
if val is None:
inner_dict[key] = ''
elif isinstance(val,
str) and val.lower() in ['up', 'down']:
color = 'red' if val.lower() == 'down' else 'green'
if tablefmt == 'rich':
inner_dict[
key] = f'[b {color}]{val.title()}[/b {color}]'
else:
inner_dict[key] = typer.style(val.title(),
fg=color)
else:
if tablefmt == 'rich':
inner_dict[key] = Text(str(val), style=None)
else:
inner_dict[key] = str(val)
else:
val = self.listify(val)
if val and tablefmt == "rich" and hasattr(
val[0], 'keys'):
inner_table = Table(
*(k for k in val[0].keys()),
show_header=True,
# padding=(0, 0),
pad_edge=False,
collapse_padding=True,
show_edge=False,
header_style="bold cyan",
box=SIMPLE)
_ = [
inner_table.add_row(*[
self.do_pretty(kk, str(vv))
for kk, vv in v.items()
]) for v in val
]
with console.capture():
console.print(inner_table)
inner_dict[key] = console.export_text()
elif val and tablefmt == "tabulate" and hasattr(
val[0], 'keys'):
inner_table = tabulate(val,
headers="keys",
tablefmt=tablefmt)
inner_dict[key] = inner_table
else:
if all(isinstance(v, str) for v in val):
inner_dict[key] = ", ".join(val)
out.append(inner_dict)
return out
raw_data = outdata
_lexer = table_data = None
if config and config.sanitize and raw_data and all(
isinstance(x, dict) for x in raw_data):
redact = [
"mac", "serial", "neighborMac", "neighborSerial",
"neighborPortMac", "longitude", "latitude"
]
outdata = [{
k: d[k] if k not in redact else "--redacted--"
for k in d
} for d in raw_data]
# -- // List[str, ...] \\ -- Bypass all formatters, (config file output, etc...)
if outdata and all(isinstance(x, str) for x in outdata):
tablefmt = "strings"
# -- convert List[dict] --> Dict[dev_name: dict] for yaml/json outputs
if tablefmt in ['json', 'yaml', 'yml']:
outdata = self.listify(outdata)
if outdata and 'name' in outdata[0]:
outdata: Dict[str, Dict[str, Any]] = {
item['name']:
{k: v
for k, v in item.items() if k != 'name'}
for item in outdata
}
if tablefmt == "json":
raw_data = json.dumps(outdata, indent=4)
_lexer = lexers.JsonLexer
elif tablefmt in ["yml", "yaml"]:
raw_data = yaml.dump(outdata, sort_keys=False)
_lexer = lexers.YamlLexer
elif tablefmt == "csv":
raw_data = table_data = "\n".join([
",".join([
k if outdata.index(d) == 0 else str(v)
for k, v in d.items() if k not in CUST_KEYS
]) for d in outdata
])
elif tablefmt == "rich":
from rich.console import Console
from rich.table import Table
from rich.box import HORIZONTALS, SIMPLE
from rich.text import Text
from centralcli import constants
console = Console(record=True)
customer_id, customer_name = "", ""
# outdata = self.listify(outdata)
# -- // List[dict, ...] \\ --
if outdata and all(isinstance(x, dict) for x in outdata):
customer_id = outdata[0].get("customer_id", "")
customer_name = outdata[0].get("customer_name", "")
outdata = [{k: v
for k, v in d.items() if k not in CUST_KEYS}
for d in outdata]
table = Table(
# show_edge=False,
show_header=True,
title=title,
header_style='magenta',
show_lines=False,
box=HORIZONTALS,
row_styles=['none', 'dark_sea_green'])
fold_cols = ['description']
_min_max = {'min': 10, 'max': 30}
set_width_cols = {'name': _min_max, 'model': _min_max}
full_cols = [
'mac', 'serial', 'ip', 'public ip', 'version', 'radio',
'id'
]
for k in outdata[0].keys():
if k in fold_cols:
table.add_column(k,
overflow='fold',
max_width=115,
justify='left')
elif k in set_width_cols:
table.add_column(k,
min_width=set_width_cols[k]['min'],
max_width=set_width_cols[k]['max'],
justify='left')
elif k in full_cols:
table.add_column(k, no_wrap=True, justify='left')
else:
table.add_column(k, justify='left')
formatted = _do_subtables(outdata)
[
table.add_row(*list(in_dict.values()))
for in_dict in formatted
]
if title:
table.title = f'[italic cornflower_blue]{constants.what_to_pretty(title)}'
if account or caption:
table.caption_justify = 'left'
table.caption = '' if not account else f'[italic dark_olive_green2] Account: {account}'
if caption:
table.caption = f"[italic dark_olive_green2]{table.caption} {caption}"
data_header = f"--\n{'Customer ID:':15}{customer_id}\n{'Customer Name:':15} {customer_name}\n--\n"
with console.capture():
console.print(table)
raw_data = console.export_text(clear=False)
table_data = console.export_text(styles=True)
raw_data = f"{data_header}{raw_data}" if customer_id else f"{raw_data}"
table_data = f"{data_header}{table_data}" if customer_id else f"{table_data}"
elif tablefmt == "tabulate":
customer_id = customer_name = ""
outdata = self.listify(outdata)
# -- // List[dict, ...] \\ --
if outdata and all(isinstance(x, dict) for x in outdata):
customer_id = outdata[0].get("customer_id", "")
customer_name = outdata[0].get("customer_name", "")
outdata = [{k: v
for k, v in d.items() if k not in CUST_KEYS}
for d in outdata]
raw_data = outdata
outdata = _do_subtables(outdata, tablefmt=tablefmt)
# outdata = [dict((k, v) for k, v in zip(outdata[0].keys(), val)) for val in outdata]
table_data = tabulate(outdata,
headers="keys",
tablefmt=tablefmt)
td = table_data.splitlines(keepends=True)
table_data = f"{typer.style(td[0], fg='cyan')}{''.join(td[1:])}"
data_header = f"--\n{'Customer ID:':15}{customer_id}\n" \
f"{'Customer Name:':15} {customer_name}\n--\n"
table_data = f"{data_header}{table_data}" if customer_id else f"{table_data}"
raw_data = f"{data_header}{raw_data}" if customer_id else f"{raw_data}"
else: # strings output No formatting
# -- // List[str, ...] \\ --
if len(outdata) == 1:
if "\n" not in outdata[0]:
# we can format green as only success output is sent through formatter.
table_data = typer.style(f" {outdata[0]}", fg="green")
raw_data = outdata[0]
else: # template / config file output
# get rid of double nl @ EoF (configs)
raw_data = table_data = "{}\n".format(
'\n'.join(outdata).rstrip('\n'))
else:
raw_data = table_data = '\n'.join(outdata)
# Not sure what hit's this, but it was created so something must
log.debug("List[str] else hit")
if _lexer and raw_data:
table_data = highlight(
bytes(raw_data, 'UTF-8'), _lexer(),
formatters.Terminal256Formatter(style='solarized-dark'))
return self.Output(rawdata=raw_data,
prettydata=table_data,
config=config)
# myfile.write(text+' # '+historyName+'\n')
else:
print("ERR: entry not found")
main()
from rich.panel import Panel
from rich.table import Table
from rich import box
from prompt_toolkit import print_formatted_text, ANSI
console = Console()
with console.capture() as capture:
grid = Table(expand=False,
box=box.SQUARE,
show_header=False,
show_edge=False)
text = """ ____ __
/ __// /_ [#444444 on blue]v0.12[/#444444 on blue]
_\ \ / __/ [#444444 on green]githubThoms10[/#444444 on green]
/___/ \__/"""
text2 = """███████ ████████
██ ██
███████ ██
██ ██
███████ ██"""
#grid.add_column(style="red")
#grid.add_column(style="green on purple")
def capture_show_output(instance):
console = Console()
with console.capture() as capture:
instance.info(console)
return capture.get()
def render_text(self) -> str:
"""Exports the info to string"""
console = Console(record=True)
with console.capture():
self.info(console)
return console.export_text()
def test_describe(monkeypatch):
monkeypatch.setenv(
"MODELKIT_ASSETS_DIR", os.path.join(TEST_DIR, "testdata", "test-bucket")
)
class SomeSimpleValidatedModelWithAsset(Model[str, str]):
"""
This is a summary
that also has plenty more text
"""
CONFIGURATIONS = {"some_model_a": {"asset": "assets-prefix"}}
def _predict(self, item):
return item
class SomeSimpleValidatedModelA(Model[str, str]):
"""
This is a summary
that also has plenty more text
"""
CONFIGURATIONS = {"some_model_a": {}}
def _predict(self, item):
return item
class ItemModel(pydantic.BaseModel):
string: str
class ResultModel(pydantic.BaseModel):
sorted: str
class A:
def __init__(self):
self.x = 1
self.y = 2
class SomeComplexValidatedModelA(Model[ItemModel, ResultModel]):
"""
More complex
With **a lot** of documentation
"""
CONFIGURATIONS = {
"some_complex_model_a": {
"model_dependencies": ["some_model_a"],
"asset": os.path.join(
TEST_DIR,
"testdata",
"test-bucket",
"assets-prefix",
"category",
"asset",
"0.0",
),
"model_settings": {"batch_size": 128},
}
}
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.some_object = A()
def _predict(self, item):
return item
# test without a console and no models
library = ModelLibrary()
library.describe()
# test with assets
library = ModelLibrary(
models=[
SomeSimpleValidatedModelA,
SomeSimpleValidatedModelWithAsset,
SomeComplexValidatedModelA,
]
)
library.describe()
# test with models but not assets
library = ModelLibrary(
models=[SomeSimpleValidatedModelA, SomeComplexValidatedModelA]
)
console = Console()
with console.capture() as capture:
library.describe(console=console)
if platform.system() != "Windows":
# Output is different on Windows platforms since
# modelkit.utils.memory cannot track memory increment
# and write it
r = ReferenceText(os.path.join(TEST_DIR, "testdata"))
captured = capture.get()
EXCLUDED = ["load time", "load memory", "asset", "category/asset", os.path.sep]
captured = "\n".join(
line
for line in captured.split("\n")
if not any(x in line for x in EXCLUDED)
)
r.assert_equal("library_describe.txt", captured)
def __init__(
self,
# ModelLibrary arguments
required_models: Optional[List[str]] = None,
settings: Optional[Union[Dict, LibrarySettings]] = None,
assetsmanager_settings: Optional[dict] = None,
configuration: Optional[Dict[str, Union[Dict[str, Any],
ModelConfiguration]]] = None,
models: Optional[LibraryModelsType] = None,
# paths overrides change the configuration key into a path
route_paths: Optional[Dict[str, str]] = None,
# APIRouter arguments
**kwargs,
) -> None:
super().__init__(
required_models=required_models,
settings=settings,
assetsmanager_settings=assetsmanager_settings,
configuration=configuration,
models=models,
**kwargs,
)
route_paths = route_paths or {}
for model_name in self.lib.required_models:
m: AbstractModel = self.lib.get(model_name)
if not isinstance(m, AbstractModel):
continue
path = route_paths.get(model_name, "/predict/" + model_name)
batch_path = route_paths.get(model_name,
"/predict/batch/" + model_name)
summary = ""
description = ""
if m.__doc__:
doclines = m.__doc__.strip().split("\n")
summary = doclines[0]
if len(doclines) > 1:
description = "".join(doclines[1:])
console = Console(no_color=True)
with console.capture() as capture:
t = m.describe()
console.print(t)
description += "\n\n```" + str(capture.get()) + "```"
logger.info("Adding model", name=model_name)
item_type = m._item_type or Any
try:
item_type.schema() # type: ignore
except (ValueError, AttributeError):
logger.info("Discarding item type info for model",
name=model_name,
path=path)
item_type = Any
self.add_api_route(
path,
self._make_model_endpoint_fn(m, item_type),
methods=["POST"],
description=description,
summary=summary,
tags=[str(type(m).__module__)],
)
self.add_api_route(
batch_path,
self._make_batch_model_endpoint_fn(m, item_type),
methods=["POST"],
description=description,
summary=summary,
tags=[str(type(m).__module__)],
)
logger.info("Added model to service", name=model_name, path=path)
