Coverage for pydantic_ai_jupyter / decorators.py: 44%

1from typing import Any, Callable, Literal, TypeGuard, TypeVar 

2 

3from .protocols import HTMLRepresentable, MarkdownRepresentable 

4 

5 

6def is_html_representable(obj: Any) -> TypeGuard[HTMLRepresentable]: 

7 return hasattr(obj, "to_html") and callable(obj.to_html) 

8 

9 

10def is_markdown_representable(obj: Any) -> TypeGuard[MarkdownRepresentable]: 

11 return hasattr(obj, "to_markdown") and callable(obj.to_markdown) 

12 

13 

14T = TypeVar("T") 

15 

16 

17def renderable( 

18 cls: type[T] | None = None, default: Literal["markdown", "html"] = "html" 

19) -> type[T] | Callable[[type[T]], type[T]]: 

20 """Decorate a class to make `render` functions the way to display in IPython""" 

21 

22 def decorator(target_cls: type[T]) -> type[T]: 

23 def _repr_mimebundle_( 

24 self, include=None, exclude=None 

25 ) -> dict[str, str | dict]: 

26 # Allow the user to pass back a string of HTML, a VDOM object, or other displayable 

27 rendered = self.render() 

28 

29 if isinstance(rendered, str): 

30 if default == "html": 

31 return {"text/html": rendered} 

32 else: 

33 return {"text/markdown": rendered} 

34 elif is_markdown_representable(rendered): 

35 return {"text/markdown": rendered.to_markdown()} 

36 elif is_html_representable(rendered): 

37 return {"text/html": rendered.to_html()} 

38 else: 

39 from IPython import get_ipython 

40 

41 ip = get_ipython() 

42 if ip is None or ip.display_formatter is None: 

43 raise ValueError("render() must return a string or a VDOM object") 

44 mime_bundle = ip.display_formatter.format(rendered) 

45 return mime_bundle 

46 

47 raise ValueError("render() must return something displayable") 

48 

49 setattr(target_cls, "_repr_mimebundle_", _repr_mimebundle_) 

50 return target_cls 

51 

52 if cls is None: 

53 # Called with arguments: @renderable(default="markdown") 

54 return decorator 

55 else: 

56 # Called without arguments: @renderable 

57 return decorator(cls) 

58 

59 

60def markdown(cls: type[T]) -> type[T]: 

61 """Decorate a class to make `render` functions return rendered Markdown in Jupyter""" 

62 return renderable(cls, default="markdown") 

63 

64 

65def html(cls: type[T]) -> type[T]: 

66 """Decorate a class to make `render` functions return rendered HTML in Jupyter""" 

67 return renderable(cls, default="html")