Question

是否可以向熊猫CategoricalDtype添加个性化的方法和属性？我应该使用类继承还是类似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工作，则需要以某种方式告诉该系列使用您创建的子类访问器，而不是默认的访问器。我不知道该怎么做，尽管有人可以。您可以随后在其访问器上添加猴子修补程序，但是每次创建新系列时都必须这样做，因此这看起来非常脆弱且无法维护。不过，您可以做的是使用 separate 自定义访问器，而不要通过.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)会引发同样的错误。

有一个类似的功能可以为数据帧注册访问器。

向CategoricalDtype添加个性化的方法和属性

1 个答案: