我想把pandas的自定义访问器与定义原生属性这两者结合起来。

我希望把一个原生属性定义在访问器中，并能在管道操作中继续生效。

请看下面的示例：

import pandas as pd

@pd.api.extensions.register_dataframe_accessor("geo")
class GeoAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    def start(self):
        print("start")
        self._initial_shape = self._obj.shape
        return self._obj

    def end(self):
        print("Previous shape:", self._initial_shape)
        return self._obj

df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, None]})
df.geo.start().dropna(subset=["y"]).geo.end()

这会返回一个错误：

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[11], line 2
      1 df=pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, None]})
----> 2 df.geo.start().dropna(subset=["y"]).geo.end()

Cell In[9], line 13
     12 def end(self):
---> 13     print("Previous shape:", self._initial_shape)
     14     return self._obj

AttributeError: 'GeoAccessor' object has no attribute '_initial_shape'

不幸的是，在调用 end() 时，属性 _initial_shape 不可用。

解决方案

我通过在DataFrame内部创建一个自定义属性，并在 _metadata 列表中指定属性名，来实现这一点：

import pandas as pd

@pd.api.extensions.register_dataframe_accessor("geo")
class GeoAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj
        self._obj._metadata += ["_initial_shape"]

    def start(self):
        print("Initial shape:", self._obj.shape)
        self._obj._initial_shape = self._obj.shape
        return self._obj

    def end(self):
        print("Previous shape:", self._obj._initial_shape)
        return self._obj

df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, None]})
df.geo.start().dropna(subset=["y"]).geo.end()

返回

Initial shape: (3, 2)
Previous shape: (3, 2)
[The dataframe]