TYP: register_*_accessor decorators (#58339)

Author: twoertwein

.pre-commit-config.yaml

@@ -108,7 +108,7 @@ repos:
         types: [python]
         stages: [manual]
         additional_dependencies: &pyright_dependencies
-        - pyright@1.1.351
+        - pyright@1.1.352
     -   id: pyright
         # note: assumes python env is setup and activated
         name: pyright reportGeneralTypeIssues

doc/source/whatsnew/v3.0.0.rst

@@ -142,6 +142,8 @@ Optional libraries below the lowest tested version may still work, but are not c
 +------------------------+---------------------+
 | adbc-driver-postgresql | 0.10.0              |
 +------------------------+---------------------+
+| mypy (dev)             | 1.9.0               |
++------------------------+---------------------+
 
 See :ref:`install.dependencies` and :ref:`install.optional_dependencies` for more.
 

environment.yml

@@ -77,7 +77,7 @@ dependencies:
 
   # code checks
   - flake8=6.1.0  # run in subprocess over docstring examples
-  - mypy=1.8.0  # pre-commit uses locally installed mypy
+  - mypy=1.9.0  # pre-commit uses locally installed mypy
   - tokenize-rt  # scripts/check_for_inconsistent_pandas_namespace.py
   - pre-commit>=3.6.0
 

pandas/_typing.py

@@ -239,6 +239,7 @@ def __reversed__(self) -> Iterator[_T_co]: ...
 # see https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators
 FuncType = Callable[..., Any]
 F = TypeVar("F", bound=FuncType)
+TypeT = TypeVar("TypeT", bound=type)
 
 # types of vectorized key functions for DataFrame::sort_values and
 # DataFrame::sort_index, among others

pandas/core/accessor.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 from typing import (
+    TYPE_CHECKING,
     Callable,
     final,
 )
@@ -16,6 +17,12 @@
 from pandas.util._decorators import doc
 from pandas.util._exceptions import find_stack_level
 
+if TYPE_CHECKING:
+    from pandas._typing import TypeT
+
+    from pandas import Index
+    from pandas.core.generic import NDFrame
+
 
 class DirNamesMixin:
     _accessors: set[str] = set()
@@ -232,7 +239,9 @@ def __get__(self, obj, cls):
 
 
 @doc(klass="", examples="", others="")
-def _register_accessor(name: str, cls):
+def _register_accessor(
+    name: str, cls: type[NDFrame | Index]
+) -> Callable[[TypeT], TypeT]:
     """
     Register a custom accessor on {klass} objects.
 
@@ -277,7 +286,7 @@ def _register_accessor(name: str, cls):
     {examples}
     """
 
-    def decorator(accessor):
+    def decorator(accessor: TypeT) -> TypeT:
         if hasattr(cls, name):
             warnings.warn(
                 f"registration of accessor {accessor!r} under name "
@@ -320,7 +329,7 @@ def decorator(accessor):
 
 
 @doc(_register_accessor, klass="DataFrame", examples=_register_df_examples)
-def register_dataframe_accessor(name: str):
+def register_dataframe_accessor(name: str) -> Callable[[TypeT], TypeT]:
     from pandas import DataFrame
 
     return _register_accessor(name, DataFrame)
@@ -351,7 +360,7 @@ def register_dataframe_accessor(name: str):
 
 
 @doc(_register_accessor, klass="Series", examples=_register_series_examples)
-def register_series_accessor(name: str):
+def register_series_accessor(name: str) -> Callable[[TypeT], TypeT]:
     from pandas import Series
 
     return _register_accessor(name, Series)
@@ -385,7 +394,7 @@ def register_series_accessor(name: str):
 
 
 @doc(_register_accessor, klass="Index", examples=_register_index_examples)
-def register_index_accessor(name: str):
+def register_index_accessor(name: str) -> Callable[[TypeT], TypeT]:
     from pandas import Index
 
     return _register_accessor(name, Index)

requirements-dev.txt

@@ -54,7 +54,7 @@ moto
 flask
 asv>=0.6.1
 flake8==6.1.0
-mypy==1.8.0
+mypy==1.9.0
 tokenize-rt
 pre-commit>=3.6.0
 gitpython