cli.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. import contextlib
  2. import logging
  3. from typing import Any, Dict, Generator, List, Optional, Tuple
  4. import typer
  5. from httpx import HTTPError, HTTPStatusError, ReadTimeout
  6. from rich.segment import Segment
  7. from rich_toolkit import RichToolkit, RichToolkitTheme
  8. from rich_toolkit.progress import Progress
  9. from rich_toolkit.styles import MinimalStyle, TaggedStyle
  10. logger = logging.getLogger(__name__)
  11. class FastAPIStyle(TaggedStyle):
  12. def __init__(self, tag_width: int = 11):
  13. super().__init__(tag_width=tag_width)
  14. def _get_tag_segments(
  15. self,
  16. metadata: Dict[str, Any],
  17. is_animated: bool = False,
  18. done: bool = False,
  19. ) -> Tuple[List[Segment], int]:
  20. if not is_animated:
  21. return super()._get_tag_segments(metadata, is_animated, done)
  22. emojis = [
  23. "🥚",
  24. "🐣",
  25. "🐤",
  26. "🐥",
  27. "🐓",
  28. "🐔",
  29. ]
  30. tag = emojis[self.animation_counter % len(emojis)]
  31. if done:
  32. tag = emojis[-1]
  33. left_padding = self.tag_width - 1
  34. left_padding = max(0, left_padding)
  35. return [Segment(tag)], left_padding
  36. def get_rich_toolkit(minimal: bool = False) -> RichToolkit:
  37. style = MinimalStyle() if minimal else FastAPIStyle(tag_width=11)
  38. theme = RichToolkitTheme(
  39. style=style,
  40. theme={
  41. "tag.title": "white on #009485",
  42. "tag": "white on #007166",
  43. "placeholder": "grey85",
  44. "text": "white",
  45. "selected": "#007166",
  46. "result": "grey85",
  47. "progress": "on #007166",
  48. "error": "red",
  49. },
  50. )
  51. return RichToolkit(theme=theme)
  52. @contextlib.contextmanager
  53. def handle_http_errors(
  54. progress: Progress,
  55. message: Optional[str] = None,
  56. ) -> Generator[None, None, None]:
  57. try:
  58. yield
  59. except ReadTimeout as e:
  60. logger.debug(e)
  61. progress.set_error(
  62. "The request to the FastAPI Cloud server timed out. Please try again later."
  63. )
  64. raise typer.Exit(1) from None
  65. except HTTPError as e:
  66. logger.debug(e)
  67. # Handle validation errors from Pydantic models, this should make it easier to debug :)
  68. if isinstance(e, HTTPStatusError) and e.response.status_code == 422:
  69. logger.debug(e.response.json()) # pragma: no cover
  70. if isinstance(e, HTTPStatusError) and e.response.status_code in (401, 403):
  71. message = "The specified token is not valid. Use `fastapi login` to generate a new token."
  72. else:
  73. message = (
  74. message
  75. or f"Something went wrong while contacting the FastAPI Cloud server. Please try again later. \n\n{e}"
  76. )
  77. progress.set_error(message)
  78. raise typer.Exit(1) from None