How can I change a class to force it to return a pandas dataframe, not as a object but in a similar fashion to how a function would.
In other words, turn this class that gives me the error AssertionError: DataFrame Expected type <class 'pandas.core.frame.DataFrame'>, found <class 'NoneType'> instead :
minimum example: class that is meant to return dataframe
import pandas as pd
from pandas.testing import assert_frame_equal
class SumAllClassTwo:
''' class sums class_2 to class_2.5 columns then writes it to class_2'''
def __init__(self, df:pd.DataFrame):
self.df = df
def sum_all_class_two(self):
self.df['class_2'] = self.df['class_2'] + self.df['class_2.5'] # sum 2 and 2.5
self.df = self.df.drop(columns=['class_2.5']) # return new 2
return self.df
class TestSumAllClassTwo:
def test_one():
input_df = pd.DataFrame({'class_2':[1], 'class_2.5':[1]})
result = SumAllClassTwo(input_df) # don't want df as an attribute
expected = pd.DataFrame({'class_2':[2]})
assert_frame_equal(result, expected)
# run test class
TestSumAllClassTwo.test_one()
so that it returns a dataframe as the similar method, which is working as intended, would:
working method equivalent
def sum_all_class_two_2(df:pd.DataFrame=None):
df['class_2'] = df['class_2'] + df['class_2.5'] # sum 2 and 2.5
df = df.drop(columns=['class_2.5']) # return new 2
return df
def test_sum_all_class_two_2():
expected = pd.DataFrame({'class_2':[2]})
input_df = pd.DataFrame({'class_2':[1], 'class_2.5':[1]})
result = sum_all_class_two_2(input_df)
assert_frame_equal(expected, result)
returning a 2 in class 2 by summing the values in row 1 of class_2 and class_2.5.