主页 / coding / 问题 / 78501868
Accepted
Bogaso
Asked: 2024-05-19 15:25:17 +0800 CST

创建自定义函数时定义数据类型

  • 772

我有以下自定义功能

import pandas as pd
def MyFn(DF : pd.DataFrame) -> float :
  return DF['Col_A'].values[1] - DF['Col_B'].values[1]

但是我想强制用户提供一个包含 2 列且列名称为的数据框'Col_A' and 'Col_B'

任何关于我如何做到这一点的见解将不胜感激。

python

2 个回答

  1. 这是我的方法

    我使用 try 函数来防止异常，这是一种简单的修复方法，而且更灵活

    import pandas as pd
def MyFn(DF : pd.DataFrame) -> float :
  try:
    return DF['Col_A'].values[1] - DF['Col_B'].values[1]
  except:
    print("you must include 'Col_A' and 'Col_B' in the argument!")
    return 0.0 # I would like to return something to prevent error

    希望这能解决问题。

    • 2
  2. Best Answer

    您应该能够使用assertand DF.columns.tolist() == ['Col_A', 'Col_B']：

    def MyFn(DF : pd.DataFrame) -> float :
    assert DF.columns.tolist() == ['Col_A', 'Col_B'], 'DF must have two columns named "Col_A" and "Col_B"'
    return DF['Col_A'].values[1] - DF['Col_B'].values[1]

MyFn(pd.DataFrame([[1, 2], [3, 5]], columns=['Col_A', 'Col_B']))
# -2

MyFn(pd.DataFrame([[1, 2], [3, 5]], columns=['Col_A', 'Col_C']))
# AssertionError: DF must have two columns named "Col_A" and "Col_B"

MyFn(pd.DataFrame([[1, 2, 3], [4, 6, 7]], columns=['Col_A', 'Col_B', 'Col_C']))
# AssertionError: DF must have two columns named "Col_A" and "Col_B"

    请注意，这将需要按此特定顺序恰好两列“Col_A”和“Col_B”。

    更灵活的选择是使用try/ except，这将允许您独立处理另一个潜在错误：传递少于两行的 DataFrame 会触发 IndexError：

    def MyFn(DF : pd.DataFrame) -> float :
    try:
        return DF['Col_A'].values[1] - DF['Col_B'].values[1]
    except KeyError as e:
        raise Exception('DF must contain "Col_A" and "Col_B"') from e
    except IndexError as e:
        raise Exception('DF must have at least two rows') from e

MyFn(pd.DataFrame([[1, 2], [3, 5]], columns=['Col_A', 'Col_B']))
# -2

#MyFn(pd.DataFrame([[1, 2], [3, 5]], columns=['Col_A', 'Col_C']))
# Exception: DF must contain "Col_A" and "Col_B"

MyFn(pd.DataFrame([[1, 2, 3], [4, 6, 7]], columns=['Col_A', 'Col_B', 'Col_C']))
# -2

MyFn(pd.DataFrame([[1, 2]], columns=['Col_A', 'Col_B']))
# Exception: DF must have at least two rows
    • 0

相关问题