Определите тип данных при создании пользовательской функции
У меня есть пользовательская функция ниже
import pandas as pd
def MyFn(DF : pd.DataFrame) -> float :
return DF['Col_A'].values[1] - DF['Col_B'].values[1]
Однако я хочу заставить пользователя предоставить фрейм данных с двумя столбцами и иметь имена столбцов как 'Col_A' and 'Col_B'
Любое понимание того, как я могу это сделать, было бы очень признательно.
Это работает. Однако есть ли способ уместить такое условие в рамки DF : pd.DataFrame?
Если вы имеете в виду использование набора текста, это не помешает вам передать неверные данные. Например, MyFn({'Col_A': pd.Series([0, 1]), 'Col_B': pd.Series([4, 5])}) не вызовет исключения, даже если мы передали словарь, а не DataFrame.
Ответы 2
Вы должны уметь использовать assert и 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
Вот мой метод
Я использую функцию 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
Надеюсь, это решит проблему.
Заставить как? Подойдет ли
assert DF.columns.tolist() == ['Col_A', 'Col_B']вам?