Pythonic类型的提示与熊猫?

时间:2017-05-10 11:13:43

标签: python typing

让我们采用一个简单的函数,它接受一个str并返回一个数据帧:

import pandas as pd
def csv_to_df(path):
    return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

为此功能添加类型提示的推荐pythonic方法是什么?

如果我向python询问DataFrame的类型,则返回pandas.core.frame.DataFrame。 以下不会起作用,因为它会告诉我大熊猫没有定义。

 def csv_to_df(path: str) -> pandas.core.frame.DataFrame:
     return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

5 个答案:

答案 0 :(得分:24)

为什么不使用pd.DataFrame

def csv_to_df(path: str) -> pd.DataFrame:
    return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

结果是一样的:

> help(csv_to_df)
Help on function csv_to_df in module __main__:
csv_to_df(path:str) -> pandas.core.frame.DataFrame

答案 1 :(得分:8)

我目前正在做以下事情:

from typing import TypeVar
PandasDataFrame = TypeVar('pandas.core.frame.DataFrame')
def csv_to_df(path: str) -> PandasDataFrame:
    return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

给出了:

> help(csv_to_df)
Help on function csv_to_df in module __main__:

csv_to_df(path:str) -> ~pandas.core.frame.DataFrame

不知道pythonic是怎样的,但我发现它可以理解为类型提示。

答案 2 :(得分:3)

这是从原始问题出发,但是使用TypeVar和@Georgy的注释建立了@dangom的答案,即无法在类型提示中为DataFrame列指定数据类型,可以使用一种简单的解决方法像这样在DataFrame中指定数据类型:

from typing import TypeVar
DataFrameStr = TypeVar("pandas.core.frame.DataFrame(str)")
def csv_to_df(path: str) -> DataFrameStr:
    return pd.read_csv(path, skiprows=1, sep='\t', comment='#')

答案 3 :(得分:0)

现在有一个pip软件包可以帮助您解决这个问题。 https://github.com/CedricFR/dataenforce

您可以使用pip install dataenforce进行安装,并使用类似pythonic的类型提示:

def preprocess(dataset: Dataset["id", "name", "location"]) -> Dataset["location", "count"]:
    pass

答案 4 :(得分:0)

查看给出的答案here,该答案说明了软件包data-science-types的用法。

pip install data-science-types

演示

# program.py

import pandas as pd

df: pd.DataFrame = pd.DataFrame({'col1': [1,2,3], 'col2': [4,5,6]}) # OK
df1: pd.DataFrame = pd.Series([1,2,3]) # error: Incompatible types in assignment

以同样的方式使用mypy运行:

$ mypy program.py