如何将 doctest 与 pandas dataframe 一起使用？

Question

我在一个名为test.py （见下文）的文件中有一个简单的 function （ log_return ），我试图用一个 doctest 来实现。

import pandas as pd

def log_return(df):
    '''Return the log return based on closing prices
    
    >>> df = pd.DataFrame({'Close': [100, 101, 102, 99]}, index = pd.date_range('2022-01-01', periods=4, freq='D'))
    >>> log_return(df)

            Close         r
2022-01-01    100       NaN
2022-01-02    101  0.009950
2022-01-03    102  0.009852
2022-01-04     99 -0.029853
    '''
    df['r'] = np.log(df['Close']).diff()

但是，当我尝试从命令行执行 doctest（例如$ python test.py ）时，出现以下与空格相关的错误。 我该如何解决这个错误？

ValueError: line 5 of the docstring for __main__.log_return has inconsistent leading whitespace: '2022-01-01    100       NaN'

Answer 1

你需要像这样缩进它：

    '''
    ...

    >>> log_return(df)
                Close         r
    2022-01-01    100       NaN
    2022-01-02    101  0.009950
    2022-01-03    102  0.009852
    2022-01-04     99 -0.029853
    '''

请记住，文档测试看起来像交互式片段，因此在这种情况下，这意味着事物应该像在交互式 session（提示和输出）中一样对齐。

现在，一旦你解决了这个问题，测试就会失败，但这是一个单独的问题。