1+ from __future__ import annotations
2+
13import glob
24import re
35from pathlib import Path
6+ from typing import TYPE_CHECKING , Any , Literal , TypedDict
47
58import logistro
69
10+ if TYPE_CHECKING :
11+ from typing_extensions import TypeGuard
12+
713_logger = logistro .getLogger (__name__ )
814
915# constants
1016DEFAULT_EXT = "png"
1117DEFAULT_SCALE = 1
1218DEFAULT_WIDTH = 700
1319DEFAULT_HEIGHT = 500
14- SUPPORTED_FORMATS = ("png" , "jpg" , "jpeg" , "webp" , "svg" , "json" , "pdf" ) # pdf and eps
20+ SUPPORTED_FORMATS = ("png" , "jpg" , "jpeg" , "webp" , "svg" , "json" , "pdf" )
21+ FormatString = Literal ["png" , "jpg" , "jpeg" , "webp" , "svg" , "json" , "pdf" ]
22+
23+
24+ def _assert_format (ext : str ) -> TypeGuard [FormatString ]:
25+ if ext not in SUPPORTED_FORMATS :
26+ raise ValueError (f"File format { ext } is not supported." )
27+ return True
28+
29+
30+ Figurish = Any # Be nice to make it more specific, dictionary or something
1531
1632
17- def _is_figurish (o ):
33+ def _is_figurish (o ) -> TypeGuard [ Figurish ] :
1834 valid = hasattr (o , "to_dict" ) or (isinstance (o , dict ) and "data" in o )
1935 if not valid :
2036 _logger .debug (
@@ -43,32 +59,49 @@ def _get_figure_dimensions(layout, width, height):
4359
4460
4561def _get_format (extension ):
46- # Normalize format
4762 original_format = extension
4863 extension = extension .lower ()
4964 if extension == "jpg" :
50- extension = "jpeg"
65+ return "jpeg"
5166
5267 if extension not in SUPPORTED_FORMATS :
53- supported_formats_str = repr (list (SUPPORTED_FORMATS ))
5468 raise ValueError (
5569 f"Invalid format '{ original_format } '.\n "
56- f" Supported formats: { supported_formats_str } " ,
70+ f" Supported formats: { SUPPORTED_FORMATS !s } " ,
5771 )
5872 return extension
5973
6074
61- def to_spec (figure , layout_opts ):
75+ # Input of to_spec
76+ class LayoutOpts (TypedDict , total = False ):
77+ format : FormatString | None
78+ scale : int | float
79+ height : int | float
80+ width : int | float
81+
82+
83+ # Output of to_spec
84+ class Spec (TypedDict ):
85+ format : FormatString
86+ width : int | float
87+ height : int | float
88+ scale : int | float
89+ data : Figurish
90+
91+
92+ def to_spec (figure , layout_opts : LayoutOpts ) -> Spec :
6293 # Get figure layout
6394 layout = figure .get ("layout" , {})
6495
6596 for k , v in layout_opts .items ():
6697 if k == "format" :
6798 if v is not None and not isinstance (v , (str )):
68- raise TypeError (f"{ v } must be string or None" )
99+ raise TypeError (
100+ f"{ k } must be one of { SUPPORTED_FORMATS !s} or None, not { v } ." ,
101+ )
69102 elif k in ("scale" , "height" , "width" ):
70103 if v is not None and not isinstance (v , (float , int )):
71- raise TypeError (f"{ v } must be numeric or None" )
104+ raise TypeError (f"{ k } must be numeric or None, not { v } . " )
72105 else :
73106 raise AttributeError (f"Unknown key in layout options, { k } " )
74107
@@ -90,7 +123,7 @@ def to_spec(figure, layout_opts):
90123 }
91124
92125
93- def _next_filename (path , prefix , ext ):
126+ def _next_filename (path , prefix , ext ) -> str :
94127 default = 1 if (path / f"{ prefix } .{ ext } " ).exists () else 0
95128 re_number = re .compile (
96129 r"^" + re .escape (prefix ) + r"\-(\d+)\." + re .escape (ext ) + r"$" ,
@@ -106,7 +139,11 @@ def _next_filename(path, prefix, ext):
106139 return f"{ prefix } .{ ext } " if n == 1 else f"{ prefix } -{ n } .{ ext } "
107140
108141
109- def build_fig_spec (fig , path , opts ): # noqa: C901
142+ def build_fig_spec ( # noqa: C901, PLR0912
143+ fig : Figurish ,
144+ path : Path | str | None ,
145+ opts : LayoutOpts | None ,
146+ ) -> tuple [Spec , Path ]:
110147 if not opts :
111148 opts = {}
112149
@@ -122,23 +159,27 @@ def build_fig_spec(fig, path, opts): # noqa: C901
122159 raise TypeError ("Path should be a string or `pathlib.Path` object (or None)" )
123160
124161 if path and path .suffix and not opts .get ("format" ):
125- opts ["format" ] = path .suffix .lstrip ("." )
162+ ext = path .suffix .lstrip ("." )
163+ if _assert_format (ext ): # not strict necessary if but helps typeguard
164+ opts ["format" ] = ext
126165
127166 spec = to_spec (fig , opts )
128167
129168 ext = spec ["format" ]
130- full_path = None
169+
170+ full_path : Path | None = None
171+ directory : Path
131172 if not path :
132- directory = Path ()
173+ directory = Path () # use current Path
133174 elif path and (not path .suffix or path .is_dir ()):
134175 if not path .is_dir ():
135- raise ValueError (f"Directories will not be created for you: { path } " )
176+ raise ValueError (f"Directory { path } not found. Please create it. " )
136177 directory = path
137178 else :
138179 full_path = path
139180 if not full_path .parent .is_dir ():
140181 raise RuntimeError (
141- f"Cannot reach path { path } . Are all directories created?" ,
182+ f"Cannot reach path { path . parent } . Are all directories created?" ,
142183 )
143184 if not full_path :
144185 _logger .debug ("Looking for title" )
0 commit comments