Get typechecking up and running

.github/workflows/ci.yaml

@@ -97,3 +97,23 @@ jobs:
       - name: Upload coverage
         uses: codecov/codecov-action@v1
         if: ${{ success() }}
+
+lint:
+    runs-on: ubuntu-latest
+
+    steps:
+
+      - name: Checkout
+        uses: actions/checkout@v2
+
+      - name: Setup Python
+        uses: actions/setup-python@v2
+
+      - name: Install tools
+        run: pip install mypy flake8
+
+      - name: Flake8
+        run: make lint
+
+      - name: Type checking
+        run: make typecheck

Makefile

@@ -8,3 +8,6 @@ unittests:
 
 lint:
 	flake8 seaborn
+
+typecheck:
+	mypy -p seaborn._core --exclude seaborn._core.orig.py

ci/utils.txt

@@ -2,3 +2,4 @@ pytest!=5.3.4
 pytest-cov
 pytest-xdist
 flake8
+mypy

seaborn/_core/data.py

@@ -5,27 +5,22 @@
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Optional, Any
-    from collections.abc import Hashable, Mapping
     from pandas import DataFrame
-    from .typing import Vector
+    from seaborn._core.typing import DataSource, VariableSpec
 
 
 class PlotData:
     """Data table with plot variable schema and mapping to original names."""
     frame: DataFrame
-    names: dict[str, Optional[str]]
-    _source: Optional[DataFrame | Mapping]
+    names: dict[str, str | None]
+    _source: DataSource
 
     def __init__(
         self,
-        data: Optional[DataFrame | Mapping],
-        variables: Optional[dict[str, Hashable | Vector]],
+        data: DataSource,
+        variables: dict[str, VariableSpec],
     ):
 
-        if variables is None:
-            variables = {}
-
         frame, names = self._assign_variables(data, variables)
 
         self.frame = frame
@@ -34,17 +29,16 @@ def __init__(
         self._source_data = data
         self._source_vars = variables
 
-    def __contains__(self, key: Hashable) -> bool:
+    def __contains__(self, key: str) -> bool:
         """Boolean check on whether a variable is defined in this dataset."""
         return key in self.frame
 
     def concat(
         self,
-        data: Optional[DataFrame | Mapping],
-        variables: Optional[dict[str, Optional[Hashable | Vector]]],
+        data: DataSource,
+        variables: dict[str, VariableSpec] | None,
     ) -> PlotData:
         """Add, replace, or drop variables and return as a new dataset."""
-
         # Inherit the original source of the upsteam data by default
         if data is None:
             data = self._source_data
@@ -75,9 +69,9 @@ def concat(
 
     def _assign_variables(
         self,
-        data: Optional[DataFrame | Mapping],
-        variables: dict[str, Optional[Hashable | Vector]]
-    ) -> tuple[DataFrame, dict[str, Optional[str]]]:
+        data: DataSource,
+        variables: dict[str, VariableSpec],
+    ) -> tuple[DataFrame, dict[str, str | None]]:
         """
         Define plot variables given long-form data and/or vector inputs.
 
@@ -104,8 +98,11 @@ def _assign_variables(
             When variables are strings that don't appear in ``data``.
 
         """
-        plot_data: dict[str, Vector] = {}
-        var_names: dict[str, Optional[str]] = {}
+        frame: DataFrame
+        names: dict[str, str | None]
+
+        plot_data = {}
+        var_names = {}
 
         # Data is optional; all variables can be defined as vectors
         if data is None:
@@ -115,10 +112,8 @@ def _assign_variables(
         # Track https://data-apis.org/ for development
 
         # Variables can also be extracted from the index of a DataFrame
-        index: dict[str, Any]
         if isinstance(data, pd.DataFrame):
-            index = data.index.to_frame().to_dict(
-                "series")  # type: ignore  # data-sci-types wrong about to_dict return
+            index = data.index.to_frame().to_dict("series")
         else:
             index = {}
 
@@ -144,9 +139,9 @@ def _assign_variables(
             if val_as_data_key:
 
                 if val in data:
-                    plot_data[key] = data[val]  # type: ignore # fails on key: Hashable
+                    plot_data[key] = data[val]
                 elif val in index:
-                    plot_data[key] = index[val]  # type: ignore # fails on key: Hashable
+                    plot_data[key] = index[val]
                 var_names[key] = str(val)
 
             elif isinstance(val, str):
@@ -174,7 +169,7 @@ def _assign_variables(
                         )
                         raise ValueError(err)
 
-                plot_data[key] = val  # type: ignore # fails on key: Hashable
+                plot_data[key] = val
 
                 # Try to infer the name of the variable
                 var_names[key] = getattr(val, "name", None)
@@ -184,7 +179,7 @@ def _assign_variables(
         frame = pd.DataFrame(plot_data)
 
         # Reduce the variables dictionary to fields with valid data
-        names: dict[str, Optional[str]] = {
+        names = {
             var: name
             for var, name in var_names.items()
             # TODO I am not sure that this is necessary any more

seaborn/_core/mappings.py

@@ -6,23 +6,24 @@
 import pandas as pd
 import matplotlib as mpl
 
-from .rules import categorical_order, variable_type
-from ..utils import get_color_cycle, remove_na
-from ..palettes import QUAL_PALETTES, color_palette
+from seaborn._core.rules import VarType, variable_type, categorical_order
+from seaborn.utils import get_color_cycle, remove_na
+from seaborn.palettes import QUAL_PALETTES, color_palette
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Optional, Literal
     from pandas import Series
     from matplotlib.colors import Colormap, Normalize
-    from matplotlib.scale import Scale  # TODO or our own ScaleWrapper
-    from .typing import PaletteSpec
+    from matplotlib.scale import Scale
+    from seaborn._core.typing import PaletteSpec
 
 
 class SemanticMapping:
     """Base class for mappings between data and visual attributes."""
 
-    def setup(self, data: Series, scale: Optional[Scale]) -> SemanticMapping:
+    levels: list  # TODO Alternately, use keys of lookup_table?
+
+    def setup(self, data: Series, scale: Scale | None) -> SemanticMapping:
         # TODO why not just implement the GroupMapping setup() here?
         raise NotImplementedError()
 
@@ -59,7 +60,7 @@ def __call__(self, x):  # TODO types; will need to overload (wheee)
 
 class GroupMapping(SemanticMapping):
     """Mapping that does not alter any visual properties of the artists."""
-    def setup(self, data: Series, scale: Optional[Scale]) -> GroupMapping:
+    def setup(self, data: Series, scale: Scale | None) -> GroupMapping:
         self.levels = categorical_order(data)
         return self
 
@@ -69,18 +70,18 @@ class HueMapping(SemanticMapping):
 
     # TODO type the important class attributes here
 
-    def __init__(self, palette: Optional[PaletteSpec] = None):
+    def __init__(self, palette: PaletteSpec = None):
 
         self._input_palette = palette
 
     def setup(
         self,
         data: Series,  # TODO generally rename Series arguments to distinguish from DF?
-        scale: Optional[Scale],
+        scale: Scale | None,  # TODO or always have a Scale?
     ) -> HueMapping:
         """Infer the type of mapping to use and define it using this vector of data."""
-        palette: Optional[PaletteSpec] = self._input_palette
-        cmap: Optional[Colormap] = None
+        palette: PaletteSpec = self._input_palette
+        cmap: Colormap | None = None
 
         norm = None if scale is None else scale.norm
         order = None if scale is None else scale.order
@@ -141,17 +142,17 @@ def setup(
     def _infer_map_type(
         self,
         scale: Scale,
-        palette: Optional[PaletteSpec],
+        palette: PaletteSpec,
         data: Series,
-    ) -> Literal["numeric", "categorical", "datetime"]:
+    ) -> VarType:
         """Determine how to implement the mapping."""
-        map_type: Optional[Literal["numeric", "categorical", "datetime"]]
+        map_type: VarType
         if scale is not None:
             return scale.type
         elif palette in QUAL_PALETTES:
-            map_type = "categorical"
+            map_type = VarType("categorical")
         elif isinstance(palette, (abc.Mapping, abc.Sequence)):
-            map_type = "categorical"
+            map_type = VarType("categorical")
         else:
             map_type = variable_type(data)
 
@@ -160,8 +161,8 @@ def _infer_map_type(
     def _setup_categorical(
         self,
         data: Series,
-        palette: Optional[PaletteSpec],
-        order: Optional[list],
+        palette: PaletteSpec,
+        order: list | None,
     ) -> tuple[list, dict]:
         """Determine colors when the hue mapping is categorical."""
         # -- Identify the order and name of the levels
@@ -202,9 +203,9 @@ def _setup_categorical(
     def _setup_numeric(
         self,
         data: Series,
-        palette: Optional[PaletteSpec],
-        norm: Optional[Normalize],
-    ) -> tuple[list, dict, Optional[Normalize], Colormap]:
+        palette: PaletteSpec,
+        norm: Normalize | None,
+    ) -> tuple[list, dict, Normalize | None, Colormap]:
         """Determine colors when the hue variable is quantitative."""
         cmap: Colormap
         if isinstance(palette, dict):