子实例化时如何在父类之间切换？

Question

我正在尝试创建一个自定义类

CsvFrame

，这是一个用 pandas 或 Polars 制作的数据框。

为此，我编写了以下代码：

class CsvFrame:
    def __init__(self, engine, *args, **kwargs):
        if engine == 'polars':
            import polars as pl
            pl.DataFrame.__init__(pl.read_csv(*args, **kwargs))

        if engine == 'pandas':
            import pandas as pd
            pd.DataFrame.__init__(pd.read_csv(*args, **kwargs))

现在当我实例化一个对象时，有两个问题：

• VS Code-Jupyter 中没有数据框的 HTML 表示
• 数据框的任何方法或属性都不可用

import io

input_text = '''
col1,col2
A,1
B,2
'''

cfr = CsvFrame('polars', io.StringIO(input_text))

# problem 1
cfr # <__main__.CsvFrame at 0x1fd721f32c0>

# problem 2
cfr.melt()
AttributeError: 'CsvFrame' object has no attribute 'melt'

你能帮我解决这个问题吗？

Answer 1

您可以使用工厂函数，根据给定的引擎返回

polars.DataFrame

或

pandas.DataFrame

的子类实例：

def get_csv_frame(engine, *args, **kwargs):
    if engine == 'polars':
        from polars import DataFrame, read_csv
    elif engine == 'pandas':
        from pandas import DataFrame, read_csv
    else:
        raise ValueError(f'Unsupported engine {engine}')

    class CsvFrame(DataFrame):
        def __new__(cls, *args, **kwargs):
            return read_csv(*args, **kwargs)
        # more methods can be added here

    return CsvFrame(*args, **kwargs)