More type annotations

Author: mwouts

pyproject.toml

@@ -128,5 +128,15 @@ artifacts = [
 
 [tool.pyright]
 typeCheckingMode = "standard"
-ignore = ["apps/marimo"]
-strict = ["docs", "apps"]
+ignore = [
+    "apps/marimo",
+    "src/itables_for_dash/ITable.py",
+]
+strict = [
+    "docs",
+    "apps",
+    "src/itables_for_dash",
+    "src/itables/__init__.py",
+    "src/itables/dash.py",
+    "src/itables/streamlit.py",
+    ]

src/itables/datatables_format.py

@@ -1,18 +1,22 @@
 import json
 import warnings
-from typing import Optional
+from typing import Any, Optional, Sequence, Union
 
 import numpy as np
 import pandas as pd
 import pandas.io.formats.format as fmt
 
+from .typing import DataFrameOrSeries
+
 try:
     import polars as pl
 except ImportError:
     pl = None
 
 
-def _format_column(x, escape_html: bool):
+def _format_column(
+    x: "pd.Series[Any]", escape_html: bool
+) -> "Union[pd.Series[Any],Sequence[Any]]":
     dtype_kind = x.dtype.kind
     if dtype_kind in ["b", "i"]:
         return x
@@ -28,21 +32,23 @@ def _format_column(x, escape_html: bool):
         # Older versions of Pandas don't have 'leading_space'
         x = fmt.format_array(x._values, None, justify="all")  # type: ignore
 
+    y: "Union[pd.Series[Any], Sequence[Any]]" = x
     if dtype_kind == "f":
         try:
-            x = np.array(x).astype(float)
+            z = np.array(x).astype(float)
         except ValueError:
+            z = x
             pass
 
-        x = [escape_non_finite_float(f) for f in x]
+        y = [escape_non_finite_float(f) for f in z]
 
     if escape_html:
-        return [escape_html_chars(i) for i in x]
+        return [escape_html_chars(i) for i in y]
 
-    return x
+    return y
 
 
-def escape_non_finite_float(value):
+def escape_non_finite_float(value: Any) -> Any:
     """Encode non-finite float values to strings that will be parsed by parseJSON"""
     if not isinstance(value, float):
         return value
@@ -55,7 +61,7 @@ def escape_non_finite_float(value):
     return value
 
 
-def escape_html_chars(value):
+def escape_html_chars(value: Any) -> Any:
     """Escape HTML special characters"""
     if isinstance(value, str):
         from pandas.io.formats.printing import pprint_thing  # type: ignore
@@ -66,7 +72,9 @@ def escape_html_chars(value):
     return value
 
 
-def generate_encoder(warn_on_unexpected_types=True):
+def generate_encoder(warn_on_unexpected_types: bool = True) -> Any:
+    """Generate a JSON encoder that can handle special types like numpy"""
+
     class TableValuesEncoder(json.JSONEncoder):
         def default(self, o):
             if isinstance(o, (bool, int, float, str)):
@@ -97,11 +105,11 @@ def default(self, o):
 
 
 def datatables_rows(
-    df,
+    df: DataFrameOrSeries,
     column_count: Optional[int] = None,
     warn_on_unexpected_types: bool = False,
     escape_html: bool = True,
-):
+) -> str:
     """Format the values in the table and return the data, row by row, as requested by DataTables"""
     # We iterate over columns using an index rather than the column name
     # to avoid an issue in case of duplicated column names #89

src/itables/downsample.py

@@ -33,8 +33,11 @@ def as_nbytes(mem: Union[int, float, str]) -> int:
 
 
 def downsample(
-    df, max_rows: int = 0, max_columns: int = 0, max_bytes: Union[int, str] = 0
-):
+    df: DataFrameOrSeries,
+    max_rows: int = 0,
+    max_columns: int = 0,
+    max_bytes: Union[int, str] = 0,
+) -> tuple[DataFrameOrSeries, str]:
     """Return a subset of the dataframe that fits the limits"""
     org_rows, org_columns, org_bytes = len(df), len(df.columns), nbytes(df)
     max_bytes_numeric = as_nbytes(max_bytes)
@@ -67,8 +70,9 @@ def downsample(
 
 
 def shrink_towards_target_aspect_ratio(
-    rows, columns, shrink_factor, target_aspect_ratio
-):
+    rows: int, columns: int, shrink_factor: float, target_aspect_ratio: float
+) -> tuple[int, int]:
+    """Return the number of rows and columns of the shrinked dataframe"""
     # current and target aspect ratio
     aspect_ratio = rows / float(columns)
 

src/itables/javascript.py

@@ -135,7 +135,7 @@ def init_notebook_mode(
     all_interactive: bool = True,
     connected: bool = GOOGLE_COLAB,
     dt_bundle: Optional[Union[Path, str]] = None,
-):
+) -> None:
     """Load the DataTables library and the corresponding css (if connected=False),
     and (if all_interactive=True), activate the DataTables representation for all the Pandas DataFrames and Series.
 
@@ -189,7 +189,8 @@ def init_notebook_mode(
         display(HTML(generate_init_offline_itables_html(dt_bundle)))
 
 
-def get_animated_logo(display_logo_when_loading):
+def get_animated_logo(display_logo_when_loading: bool) -> str:
+    """Return the HTML for the loading logo of ITables"""
     if not display_logo_when_loading:
         return ""
     return f"<a href=https://mwouts.github.io/itables/>{read_package_file('logo/loading.svg')}</a>"
@@ -265,15 +266,15 @@ def _flat_header(df, show_index):
     return header
 
 
-def _tfoot_from_thead(thead):
+def _tfoot_from_thead(thead: str) -> str:
     header_rows = thead.split("</tr>")
     last_row = header_rows[-1]
     assert not last_row.strip(), last_row
     header_rows = header_rows[:-1]
     return "".join(row + "</tr>" for row in header_rows[::-1] if "<tr" in row) + "\n"
 
 
-def get_keys_to_be_evaluated(data) -> list[list[Union[int, str]]]:
+def get_keys_to_be_evaluated(data: Any) -> list[list[Union[int, str]]]:
     """
     This function returns the keys that need to be evaluated
     in the ITable arguments
@@ -296,7 +297,7 @@ def get_keys_to_be_evaluated(data) -> list[list[Union[int, str]]]:
     return keys_to_be_evaluated
 
 
-def replace_value(template, pattern, value):
+def replace_value(template: str, pattern: str, value: str) -> str:
     """Set the given pattern to the desired value in the template,
     after making sure that the pattern is found exactly once."""
     count = template.count(pattern)
@@ -311,15 +312,15 @@ def replace_value(template, pattern, value):
     return template.replace(pattern, value)
 
 
-def _datatables_repr_(df):
+def _datatables_repr_(df: DataFrameOrSeries) -> str:
     return to_html_datatable(df, connected=_CONNECTED)
 
 
 def to_html_datatable(
     df: DataFrameOrSeries,
     caption: Optional[str] = None,
     **kwargs: Unpack[ITableOptions],
-):
+) -> str:
     """
     Return the HTML representation of the given
     dataframe as an interactive datatable
@@ -548,8 +549,14 @@ def get_itables_extension_arguments(
 
 
 def warn_if_selected_rows_are_not_visible(
-    selected_rows, full_row_count, data_row_count, warn_on_selected_rows_not_rendered
-):
+    selected_rows: Optional[Sequence[int]],
+    full_row_count: int,
+    data_row_count: int,
+    warn_on_selected_rows_not_rendered: bool,
+) -> Sequence[int]:
+    """
+    Issue a warning if the selected rows are not within the range of rendered rows.
+    """
     if selected_rows is None:
         return []
 
@@ -589,7 +596,11 @@ def warn_if_selected_rows_are_not_visible(
     return [i for i in selected_rows if i < bottom_limit or i >= top_limit]
 
 
-def check_table_id(table_id: Optional[str], kwargs, df=None) -> str:
+def check_table_id(
+    table_id: Optional[str],
+    kwargs: Union[ITableOptions, DTForITablesOptions],
+    df: Optional[DataFrameOrSeries] = None,
+) -> str:
     """Make sure that the table_id is a valid HTML id.
 
     See also https://stackoverflow.com/questions/70579/html-valid-id-attribute-values
@@ -614,7 +625,9 @@ def check_table_id(table_id: Optional[str], kwargs, df=None) -> str:
     return table_id
 
 
-def set_default_options(kwargs: ITableOptions, *, use_to_html: bool, app_mode: bool):
+def set_default_options(
+    kwargs: ITableOptions, *, use_to_html: bool, app_mode: bool
+) -> None:
     if not app_mode:
         kwargs["connected"] = kwargs.get(
             "connected", ("dt_url" in kwargs) or _CONNECTED
@@ -681,7 +694,7 @@ def html_table_from_template(
     connected: bool,
     display_logo_when_loading: bool,
     kwargs: DTForITablesOptions,
-):
+) -> str:
     if "css" in kwargs:
         raise TypeError(
             "The 'css' argument has been deprecated, see the new "
@@ -800,7 +813,7 @@ def _filter_control(control, downsampling_warning):
     return None
 
 
-def safe_reset_index(df):
+def safe_reset_index(df: pd.DataFrame) -> pd.DataFrame:
     try:
         return df.reset_index()
     except ValueError:
@@ -825,6 +838,6 @@ def show(
     df: DataFrameOrSeries,
     caption: Optional[str] = None,
     **kwargs: Unpack[ITableOptions],
-):
+) -> None:
     """Render the given dataframe as an interactive datatable"""
     display(HTML(to_html_datatable(df, caption, **kwargs)))

src/itables/options.py

@@ -139,7 +139,7 @@
     )
 
 """Check that options have correct types"""
-warn_on_unexpected_option_type = (
+warn_on_unexpected_option_type: bool = (
     warn_on_undocumented_option and typing.is_typeguard_available()
 )
 if warn_on_unexpected_option_type:

src/itables/sample_dfs.py

@@ -3,7 +3,7 @@
 from datetime import datetime, timedelta
 from functools import lru_cache
 from itertools import cycle
-from typing import Any, cast
+from typing import Any, Sequence, cast
 
 import numpy as np
 import pandas as pd
@@ -28,9 +28,9 @@
     "timedelta",
 ]
 
-PANDAS_VERSION_MAJOR, PANDAS_VERSION_MINOR, _ = pd.__version__.split(".", 2)
-PANDAS_VERSION_MAJOR = int(PANDAS_VERSION_MAJOR)
-PANDAS_VERSION_MINOR = int(PANDAS_VERSION_MINOR)
+_pd_version_major, _pd_version_minor, _ = pd.__version__.split(".", 2)
+PANDAS_VERSION_MAJOR = int(_pd_version_major)
+PANDAS_VERSION_MINOR = int(_pd_version_minor)
 if PANDAS_VERSION_MAJOR == 0:
     COLUMN_TYPES = [type for type in COLUMN_TYPES if type != "boolean"]
 if PANDAS_VERSION_MAJOR == 2 and PANDAS_VERSION_MINOR == 1:
@@ -84,19 +84,20 @@ def get_countries(html: bool = True, climate_zone: bool = False) -> pd.DataFrame
     return df
 
 
-def get_population():
+def get_population() -> "pd.Series[float]":
     """A Pandas Series with the world population (from the world bank data)"""
     return pd.read_csv(find_package_file("samples/population.csv")).set_index(
         "Country"
     )["SP.POP.TOTL"]
 
 
-def get_indicators():
+def get_indicators() -> pd.DataFrame:
     """A Pandas DataFrame with a subset of the world bank indicators"""
     return pd.read_csv(find_package_file("samples/indicators.csv"))
 
 
-def get_df_complex_index():
+def get_df_complex_index() -> pd.DataFrame:
+    """A Pandas DataFrame with a complex index"""
     df = get_countries()
     df = df.reset_index().set_index(["region", "country"])
     df.columns = pd.MultiIndex.from_arrays(
@@ -116,7 +117,7 @@ def get_df_complex_index():
     return df
 
 
-def get_dict_of_test_dfs(N=100, M=100) -> dict[str, pd.DataFrame]:
+def get_dict_of_test_dfs(N: int = 100, M: int = 100) -> dict[str, pd.DataFrame]:
     NM_values = np.reshape(np.linspace(start=0.0, stop=1.0, num=N * M), (N, M))
 
     return {
@@ -281,7 +282,7 @@ def get_dict_of_test_dfs(N=100, M=100) -> dict[str, pd.DataFrame]:
     }
 
 
-def get_dict_of_polars_test_dfs(N=100, M=100) -> dict[str, Any]:
+def get_dict_of_polars_test_dfs(N: int = 100, M: int = 100) -> dict[str, Any]:
     import polars as pl
     import pyarrow as pa
 
@@ -300,7 +301,8 @@ def get_dict_of_polars_test_dfs(N=100, M=100) -> dict[str, Any]:
     return polars_dfs
 
 
-def get_dict_of_test_series():
+def get_dict_of_test_series() -> dict[str, Any]:
+    """A dictionary of Pandas Series"""
     series = {}
     for df_name, df in get_dict_of_test_dfs().items():
         if len(df.columns) > 6:
@@ -313,7 +315,8 @@ def get_dict_of_test_series():
     return series
 
 
-def get_dict_of_polars_test_series():
+def get_dict_of_polars_test_series() -> dict[str, Any]:
+    """A dictionary of Polars Series"""
     import polars as pl
     import pyarrow as pa
 
@@ -342,7 +345,8 @@ def generate_date_series():
     return pd.Series(pd.date_range("1677-09-23", "2262-04-10", freq="D"))
 
 
-def generate_random_series(rows, type):
+def generate_random_series(rows: int, type: str) -> Any:
+    """Generate a random Pandas Series of the given type and number of rows"""
     if type == "bool":
         return pd.Series(np.random.binomial(n=1, p=0.5, size=rows), dtype=bool)
     if type == "boolean":
@@ -382,15 +386,18 @@ def generate_random_series(rows, type):
     raise NotImplementedError(type)
 
 
-def generate_random_df(rows, columns, column_types=COLUMN_TYPES):
+def generate_random_df(
+    rows: int, columns: int, column_types: Sequence[str] = COLUMN_TYPES
+) -> pd.DataFrame:
     rows = int(rows)
     types = np.random.choice(column_types, size=columns)
-    columns = [
+    columns_names = [
         "Column{}OfType{}".format(col, type.title()) for col, type in enumerate(types)
     ]
 
     series = {
-        col: generate_random_series(rows, type) for col, type in zip(columns, types)
+        col: generate_random_series(rows, type)
+        for col, type in zip(columns_names, types)
     }
     index = pd.Index(range(rows))
     for x in series.values():
@@ -399,7 +406,7 @@ def generate_random_df(rows, columns, column_types=COLUMN_TYPES):
     return pd.DataFrame(series)
 
 
-def get_pandas_styler():
+def get_pandas_styler() -> Any:
     """This function returns a Pandas Styler object
 
     Cf. https://pandas.pydata.org/docs/user_guide/style.html