|
15 | 15 | from functools import wraps |
16 | 16 | from typing import Iterable |
17 | 17 |
|
| 18 | +import numpy as np |
18 | 19 | import pandas as pd |
19 | 20 | from pandas.api.types import ( |
| 21 | + is_categorical_dtype, |
20 | 22 | is_datetime64_dtype, |
21 | 23 | is_datetime64tz_dtype, |
22 | | - is_timedelta64_dtype, |
23 | 24 | is_period_dtype, |
| 25 | + is_timedelta64_dtype, |
24 | 26 | ) |
25 | 27 |
|
26 | 28 | from ...utils import adapt_mars_docstring |
27 | | -from .string_ import _string_method_to_handlers, SeriesStringMethod |
| 29 | +from .categorical import _categorical_method_to_handlers, SeriesCategoricalMethod |
28 | 30 | from .datetimes import _datetime_method_to_handlers, SeriesDatetimeMethod |
| 31 | +from .string_ import _string_method_to_handlers, SeriesStringMethod |
29 | 32 |
|
30 | 33 |
|
31 | 34 | class StringAccessor: |
@@ -262,6 +265,53 @@ def __dir__(self) -> Iterable[str]: |
262 | 265 | return list(s) |
263 | 266 |
|
264 | 267 |
|
| 268 | +class CategoricalAccessor: |
| 269 | + def __init__(self, series): |
| 270 | + if not is_categorical_dtype(series.dtype): |
| 271 | + raise AttributeError("Can only use .cat accessor with categorical values") |
| 272 | + self._series = series |
| 273 | + |
| 274 | + @property |
| 275 | + def ordered(self): |
| 276 | + return self._series.dtype.ordered |
| 277 | + |
| 278 | + @property |
| 279 | + def categories(self): |
| 280 | + return getattr(self, "_get_categories")() |
| 281 | + |
| 282 | + @classmethod |
| 283 | + def _gen_func(cls, method, is_property): |
| 284 | + def _inner(self, *args, **kwargs): |
| 285 | + op = SeriesCategoricalMethod( |
| 286 | + method=method, |
| 287 | + is_property=is_property, |
| 288 | + method_args=args, |
| 289 | + method_kwargs=kwargs, |
| 290 | + ) |
| 291 | + return op(self._series) |
| 292 | + |
| 293 | + if hasattr(pd.Series.cat, method): |
| 294 | + _inner = wraps(getattr(pd.Series.cat, method))(_inner) |
| 295 | + _inner.__doc__ = adapt_mars_docstring( |
| 296 | + getattr(pd.Series.cat, method).__doc__ |
| 297 | + ) |
| 298 | + return _inner |
| 299 | + |
| 300 | + @classmethod |
| 301 | + def _register(cls, method): |
| 302 | + # non-existing members are considered methods by default |
| 303 | + is_property = not callable(getattr(pd.Series.cat, method, lambda: None)) |
| 304 | + func = cls._gen_func(method, is_property) |
| 305 | + if is_property: |
| 306 | + func = property(func) |
| 307 | + setattr(cls, method, func) |
| 308 | + |
| 309 | + def __dir__(self) -> Iterable[str]: |
| 310 | + s = set(super().__dir__()) |
| 311 | + s.update(_categorical_method_to_handlers.keys()) |
| 312 | + return list(s) |
| 313 | + |
| 314 | + |
265 | 315 | class CachedAccessor: |
266 | 316 | def __init__(self, name: str, accessor) -> None: |
267 | 317 | self._name = name |
|
0 commit comments