向CategoricalDtyp添加个性化方法和属性

Question

有没有办法向PandasCategoricalDtype添加个性化的方法和属性？我应该使用类继承还是类似于ExtensionDtype的东西？你知道吗

例如：

vehicles = ["Plane", "Rocket", "Car", "Truck"]
vehicle_dtype = CategoricalDtype(categories=vehicles)
s = pd.Series(["Plane", "Plane", "Car"])
s = s.astype(vehicle_dtype)

有没有一个解决方案可以将方法和属性添加到vehicle_dtype来完成这样的事情？你知道吗

s.cat.is_flying
[True, True, False]

谢谢你的帮助。你知道吗

Answer 1

s.cat是pandas.core.arrays.categorical.CategoricalAccessor。如果您想s.cat.is_flying工作，您需要告诉系列使用您创建的子类访问器，而不是默认的子类访问器。我不知道该怎么做，尽管有人会这么做。您可以在以后对访问器进行monkeypatch，但每次创建新系列时都必须这样做，因此这看起来非常脆弱，不可维护。但是，您可以做的是使用一个独立的自定义访问器，而不是通过.cat。这些实际上并不难定义；请参阅文档here。下面是一个适用于您的用例的示例：

import pandas as pd

VehicleDtype = pd.api.types.CategoricalDtype(["Plane", "Rocket", "Car", "Truck"])

@pd.api.extensions.register_series_accessor("vehicle")
class VehicleAccessor:
    def __init__(self, series):
        self._validate(series)
        self._series = series

    @staticmethod
    def _validate(series):
        if not isinstance(series.dtype, CategoricalDtype) or series.dtype != VehicleDtype:
            raise TypeError("Must be VehicleDtype.")

    @property
    def is_flying(self):
        return (self._series == "Plane") | (self._series == "Rocket")

s = pd.Series(["Plane", "Plane", "Car"])
s = s.astype(VehicleDtype)

s
# 0    Plane
# 1    Plane
# 2      Car
# dtype: category
# Categories (4, object): [Plane, Rocket, Car, Truck]

s.vehicle.is_flying
# 0     True
# 1     True
# 2    False
# dtype: bool

对于类型不正确的序列，只有在尝试使用.vehicle访问器时才会抛出错误：

s2 = pd.Series(list("abcde"))  # works fine
s2.vehicle # TypeError: Must be VehicleDtype.

不过，请注意，执行dir(s2)将抛出相同的错误。你知道吗

有一个类似的函数为数据帧注册访问器。你知道吗