diff --git a/pandas/core/accessor.py b/pandas/core/accessor.py index 99b5053ce250c..aa2bf2f527bd8 100644 --- a/pandas/core/accessor.py +++ b/pandas/core/accessor.py @@ -231,7 +231,7 @@ def __get__(self, obj, cls): return accessor_obj -@doc(klass="", others="") +@doc(klass="", examples="", others="") def _register_accessor(name: str, cls): """ Register a custom accessor on {klass} objects. @@ -255,51 +255,26 @@ def _register_accessor(name: str, cls): Notes ----- - When accessed, your accessor will be initialized with the pandas object - the user is interacting with. So the signature must be + This function allows you to register a custom-defined accessor class for {klass}. + The requirements for the accessor class are as follows: - .. code-block:: python + * Must contain an init method that: - def __init__(self, pandas_object): # noqa: E999 - ... + * accepts a single {klass} object - For consistency with pandas methods, you should raise an ``AttributeError`` - if the data passed to your accessor has an incorrect dtype. + * raises an AttributeError if the {klass} object does not have correctly + matching inputs for the accessor - >>> pd.Series(["a", "b"]).dt - Traceback (most recent call last): - ... - AttributeError: Can only use .dt accessor with datetimelike values + * Must contain a method for each access pattern. - Examples - -------- - In your library code:: - - @pd.api.extensions.register_dataframe_accessor("geo") - class GeoAccessor: - def __init__(self, pandas_obj): - self._obj = pandas_obj - - @property - def center(self): - # return the geographic center point of this DataFrame - lat = self._obj.latitude - lon = self._obj.longitude - return (float(lon.mean()), float(lat.mean())) + * The methods should be able to take any argument signature. - def plot(self): - # plot this array's data on a map, e.g., using Cartopy - pass + * Accessible using the @property decorator if no additional arguments are + needed. - Back in an interactive IPython session: - - .. code-block:: ipython - - In [1]: ds = pd.DataFrame({{"longitude": np.linspace(0, 10), - ...: "latitude": np.linspace(0, 20)}}) - In [2]: ds.geo.center - Out[2]: (5.0, 10.0) - In [3]: ds.geo.plot() # plots data on a map + Examples + -------- + {examples} """ def decorator(accessor): @@ -318,21 +293,98 @@ def decorator(accessor): return decorator -@doc(_register_accessor, klass="DataFrame") +_register_df_examples = """ +An accessor that only accepts integers could +have a class defined like this: + +>>> @pd.api.extensions.register_dataframe_accessor("int_accessor") +... class IntAccessor: +... def __init__(self, pandas_obj): +... if not all(pandas_obj[col].dtype == 'int64' for col in pandas_obj.columns): +... raise AttributeError("All columns must contain integer values only") +... self._obj = pandas_obj +... +... def sum(self): +... return self._obj.sum() +... +>>> df = pd.DataFrame([[1, 2], ['x', 'y']]) +>>> df.int_accessor +Traceback (most recent call last): +... +AttributeError: All columns must contain integer values only. +>>> df = pd.DataFrame([[1, 2], [3, 4]]) +>>> df.int_accessor.sum() +0 4 +1 6 +dtype: int64""" + + +@doc(_register_accessor, klass="DataFrame", examples=_register_df_examples) def register_dataframe_accessor(name: str): from pandas import DataFrame return _register_accessor(name, DataFrame) -@doc(_register_accessor, klass="Series") +_register_series_examples = """ +An accessor that only accepts integers could +have a class defined like this: + +>>> @pd.api.extensions.register_series_accessor("int_accessor") +... class IntAccessor: +... def __init__(self, pandas_obj): +... if not pandas_obj.dtype == 'int64': +... raise AttributeError("The series must contain integer data only") +... self._obj = pandas_obj +... +... def sum(self): +... return self._obj.sum() +... +>>> df = pd.Series([1, 2, 'x']) +>>> df.int_accessor +Traceback (most recent call last): +... +AttributeError: The series must contain integer data only. +>>> df = pd.Series([1, 2, 3]) +>>> df.int_accessor.sum() +6""" + + +@doc(_register_accessor, klass="Series", examples=_register_series_examples) def register_series_accessor(name: str): from pandas import Series return _register_accessor(name, Series) -@doc(_register_accessor, klass="Index") +_register_index_examples = """ +An accessor that only accepts integers could +have a class defined like this: + +>>> @pd.api.extensions.register_index_accessor("int_accessor") +... class IntAccessor: +... def __init__(self, pandas_obj): +... if not all(isinstance(x, int) for x in pandas_obj): +... raise AttributeError("The index must only be an integer value") +... self._obj = pandas_obj +... +... def even(self): +... return [x for x in self._obj if x % 2 == 0] +>>> df = pd.DataFrame.from_dict( +... {"row1": {"1": 1, "2": "a"}, "row2": {"1": 2, "2": "b"}}, orient="index" +... ) +>>> df.index.int_accessor +Traceback (most recent call last): +... +AttributeError: The index must only be an integer value. +>>> df = pd.DataFrame( +... {"col1": [1, 2, 3, 4], "col2": ["a", "b", "c", "d"]}, index=[1, 2, 5, 8] +... ) +>>> df.index.int_accessor.even() +[2, 8]""" + + +@doc(_register_accessor, klass="Index", examples=_register_index_examples) def register_index_accessor(name: str): from pandas import Index