Update typing annotations for Python 3.9+

Author: matthewwardrop

benchmarks/plot.py

@@ -24,7 +24,7 @@ def grouped_barplot(df, cat, subcat, val, err, subcats=None, **kwargs):
             x + offsets[i],
             dfg[val].values,
             width=width,
-            label="{}".format(gr),
+            label=f"{gr}",
             yerr=dfg[err].values,
             capsize=6,
             **kwargs,

docsite/docs/guides/integration.ipynb

@@ -141,7 +141,8 @@
     }
    ],
    "source": [
-    "from typing import Iterable, List, Optional\n",
+    "from collections.abc import Iterable\n",
+    "from typing import Optional\n",
     "\n",
     "from sklearn.base import BaseEstimator, TransformerMixin\n",
     "from sklearn.linear_model import LinearRegression\n",
@@ -181,7 +182,7 @@
     "\n",
     "    def get_feature_names_out(\n",
     "        self, input_features: Optional[Iterable[str]] = None\n",
-    "    ) -> List[str]:\n",
+    "    ) -> list[str]:\n",
     "        \"\"\"\n",
     "        Expose model spec column names to scikit learn to allow column transforms later in the pipeline.\n",
     "        \"\"\"\n",

formulaic/formula.py

@@ -2,20 +2,12 @@
 
 import sys
 from abc import ABCMeta, abstractmethod
-from collections.abc import MutableSequence
+from collections.abc import Generator, Iterable, Mapping, MutableSequence
 from enum import Enum
 from typing import (
     Any,
     Callable,
-    Dict,
-    Generator,
-    Iterable,
-    List,
-    Mapping,
     Optional,
-    Set,
-    Tuple,
-    Type,
     TypeVar,
     Union,
     cast,
@@ -38,10 +30,10 @@
 FormulaSpec: TypeAlias = Union[
     "Formula",
     str,
-    List[Union[str, Term]],
-    Set[Union[str, Term]],
-    Dict[str, "FormulaSpec"],
-    Tuple["FormulaSpec", ...],
+    list[Union[str, Term]],
+    set[Union[str, Term]],
+    dict[str, "FormulaSpec"],
+    tuple["FormulaSpec", ...],
     Structured["FormulaSpec"],
 ]
 _SelfType = TypeVar("_SelfType", bound="Formula")
@@ -281,7 +273,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
-        drop_rows: Optional[Set[int]] = None,
+        drop_rows: Optional[set[int]] = None,
         **spec_overrides: Any,
     ) -> Union[ModelMatrix, Structured[ModelMatrix]]:
         """
@@ -302,7 +294,7 @@ def get_model_matrix(
 
     @property
     @abstractmethod
-    def required_variables(self) -> Set[Variable]:
+    def required_variables(self) -> set[Variable]:
         """
         The set of variables required to be in the data to materialize this
         formula.
@@ -504,7 +496,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
-        drop_rows: Optional[Set[int]] = None,
+        drop_rows: Optional[set[int]] = None,
         **spec_overrides: Any,
     ) -> Union[ModelMatrix, Structured[ModelMatrix]]:
         """
@@ -529,7 +521,7 @@ def get_model_matrix(
         )
 
     @property
-    def required_variables(self) -> Set[Variable]:
+    def required_variables(self) -> set[Variable]:
         """
         The set of variables required in the data order to materialize this
         formula.
@@ -542,7 +534,7 @@ def required_variables(self) -> Set[Variable]:
         evaluation context rather than the data context.
         """
 
-        variables: List[Variable] = [
+        variables: list[Variable] = [
             variable
             for term in self.__terms
             for factor in term.factors
@@ -605,11 +597,11 @@ def _map(
         self,
         func: Union[
             Callable[[SimpleFormula], Any],
-            Callable[[SimpleFormula, Tuple[Union[str, int], ...]], Any],
+            Callable[[SimpleFormula, tuple[Union[str, int], ...]], Any],
         ],
         recurse: bool = True,
-        as_type: Optional[Type[Structured]] = None,
-        _context: Tuple[Union[str, int], ...] = (),
+        as_type: Optional[type[Structured]] = None,
+        _context: tuple[Union[str, int], ...] = (),
     ) -> Any:
         try:
             return func(self, ())  # type: ignore
@@ -629,7 +621,7 @@ def _flatten(self) -> Generator[SimpleFormula, None, None]:
         as_of=(1, 1),
         removed_in=(2, 0),
     )
-    def _to_dict(self) -> Dict[str, SimpleFormula]:
+    def _to_dict(self) -> dict[str, SimpleFormula]:
         return {"root": self}
 
     @deprecated(
@@ -726,7 +718,7 @@ def get_model_matrix(
         self,
         data: Any,
         context: Optional[Mapping[str, Any]] = None,
-        drop_rows: Optional[Set[int]] = None,
+        drop_rows: Optional[set[int]] = None,
         **spec_overrides: Any,
     ) -> Union[ModelMatrix, Structured[ModelMatrix]]:
         """
@@ -751,7 +743,7 @@ def get_model_matrix(
         )
 
     @property
-    def required_variables(self) -> Set[Variable]:
+    def required_variables(self) -> set[Variable]:
         """
         The set of variables required in the data order to materialize this
         formula.
@@ -764,7 +756,7 @@ def required_variables(self) -> Set[Variable]:
         evaluation context rather than the data context.
         """
 
-        variables: List[Variable] = []
+        variables: list[Variable] = []
 
         # Recurse through formula to collect all variables
         self._map(
@@ -797,7 +789,7 @@ def differentiate(  # pylint: disable=redefined-builtin
         )
 
     # Ensure pickling never includes context
-    def __getstate__(self) -> Tuple[None, Dict[str, Any]]:
+    def __getstate__(self) -> tuple[None, dict[str, Any]]:
         slots = self.__slots__ + Structured.__slots__
         return (
             None,

formulaic/materializers/arrow.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Sequence
+from collections.abc import Iterator, Mapping, Sequence
+from typing import TYPE_CHECKING, Any
 
 import pandas
 from interface_meta import override
@@ -30,7 +30,7 @@ class LazyArrowTableProxy(Mapping):
     def __init__(self, table: pyarrow.Table):
         self.table = table
         self.column_names = set(self.table.column_names)
-        self._cache: Dict[str, pandas.Series] = {}
+        self._cache: dict[str, pandas.Series] = {}
         self.index = pandas.RangeIndex(len(table))
 
     def __contains__(self, value: Any) -> Any: