Skip to content

Commit 720ada5

Browse files
committed
add and fix type hints
1 parent b9e5f7a commit 720ada5

File tree

2 files changed

+46
-30
lines changed

2 files changed

+46
-30
lines changed

plotly/io/_kaleido.py

Lines changed: 43 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import json
33
from pathlib import Path
4+
from typing import Union, List
45
import importlib.metadata as importlib_metadata
56
from packaging.version import Version
67
import warnings
@@ -50,7 +51,10 @@
5051
)
5152

5253

53-
def kaleido_available():
54+
def kaleido_available() -> bool:
55+
"""
56+
Returns True if any version of Kaleido is installed, otherwise False.
57+
"""
5458
global _KALEIDO_AVAILABLE
5559
global _KALEIDO_MAJOR
5660
if _KALEIDO_AVAILABLE is not None:
@@ -64,7 +68,11 @@ def kaleido_available():
6468
return _KALEIDO_AVAILABLE
6569

6670

67-
def kaleido_major():
71+
def kaleido_major() -> int:
72+
"""
73+
Returns the major version number of Kaleido if it is installed,
74+
otherwise raises a ValueError.
75+
"""
6876
global _KALEIDO_MAJOR
6977
if _KALEIDO_MAJOR is not None:
7078
return _KALEIDO_MAJOR
@@ -146,7 +154,7 @@ def __setattr__(self, name, value):
146154
scope = None
147155

148156

149-
def as_path_object(file: str | Path) -> Path | None:
157+
def as_path_object(file: Union[str, Path]) -> Union[Path, None]:
150158
"""
151159
Cast the `file` argument, which may be either a string or a Path object,
152160
to a Path object.
@@ -185,14 +193,15 @@ def infer_format(path: Path | None, format: str | None) -> str | None:
185193

186194

187195
def to_image(
188-
fig,
189-
format=None,
190-
width=None,
191-
height=None,
192-
scale=None,
193-
validate=True,
194-
engine=None,
195-
):
196+
fig: Union[dict, plotly.graph_objects.Figure],
197+
format: Union[str, None]=None,
198+
width: Union[int, None]=None,
199+
height: Union[int, None]=None,
200+
scale: Union[int, float, None]=None,
201+
validate: bool=True,
202+
# Deprecated
203+
engine: Union[str, None]=None,
204+
) -> bytes:
196205
"""
197206
Convert a figure to a static image bytes string
198207
@@ -350,14 +359,15 @@ def to_image(
350359

351360

352361
def write_image(
353-
fig,
354-
file,
355-
format=None,
356-
scale=None,
357-
width=None,
358-
height=None,
359-
validate=True,
360-
engine="auto",
362+
fig: Union[dict, plotly.graph_objects.Figure],
363+
file: Union[str, Path],
364+
format: Union[str, None]=None,
365+
scale: Union[int, float, None]=None,
366+
width: Union[int, None]=None,
367+
height: Union[int, None]=None,
368+
validate: bool=True,
369+
# Deprecated
370+
engine: Union[str, None]="auto",
361371
):
362372
"""
363373
Convert a figure to a static image and write it to a file or writeable
@@ -481,14 +491,14 @@ def write_image(
481491

482492

483493
def write_images(
484-
fig,
485-
file,
486-
format=None,
487-
scale=None,
488-
width=None,
489-
height=None,
490-
validate=True,
491-
):
494+
fig: Union[List[Union[dict, plotly.graph_objects.Figure]], Union[dict, plotly.graph_objects.Figure]],
495+
file: Union[List[Union[str, Path]], Union[str, Path]],
496+
format: Union[List[Union[str, None]], Union[str, None]] = None,
497+
scale: Union[List[Union[int, float, None]], Union[int, float, None]] = None,
498+
width: Union[List[Union[int, None]], Union[int, None]] = None,
499+
height: Union[List[Union[int, None]], Union[int, None]] = None,
500+
validate: Union[List[bool], bool] = True,
501+
) -> None:
492502
"""
493503
Write multiple images to files or writeable objects. This is much faster than
494504
calling write_image() multiple times. This function can only be used with the Kaleido
@@ -610,7 +620,11 @@ def write_images(
610620
raise RuntimeError(PLOTLY_GET_CHROME_ERROR_MSG)
611621

612622

613-
def full_figure_for_development(fig, warn=True, as_dict=False):
623+
def full_figure_for_development(
624+
fig: Union[dict, plotly.graph_objects.Figure],
625+
warn: bool=True,
626+
as_dict: bool=False,
627+
) -> Union[plotly.graph_objects.Figure, dict]:
614628
"""
615629
Compute default values for all attributes not specified in the input figure and
616630
returns the output as a "full" figure. This function calls Plotly.js via Kaleido
@@ -677,7 +691,7 @@ def full_figure_for_development(fig, warn=True, as_dict=False):
677691
return go.Figure(fig, skip_invalid=True)
678692

679693

680-
def get_chrome():
694+
def get_chrome() -> None:
681695
"""
682696
Install Google Chrome for Kaleido
683697
This function can be run from the command line using the command `plotly_get_chrome`

plotly/io/_utils.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import List
2+
13
import plotly
24
import plotly.graph_objs as go
35
from plotly.offline import get_plotlyjs_version
@@ -43,7 +45,7 @@ def validate_coerce_output_type(output_type):
4345
return cls
4446

4547

46-
def broadcast_args_to_dicts(**kwargs):
48+
def broadcast_args_to_dicts(**kwargs: dict) -> List[dict]:
4749
"""
4850
Given one or more keyword arguments which may be either a single value or a list of values,
4951
return a list of keyword dictionaries by broadcasting the single valuesacross all the dicts.

0 commit comments

Comments
 (0)