AssertionError: DataFrame Expected type <class 'pandas.core.frame.DataFrame'>, found <class 'NoneType'> instead

Viewed 221

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.

1 Answers

Mostly not an answer but a comment with more room -

This is the constructor and doesn't return a DF it returns a class reference:

result = SumAllClassTwo(input_df)  

this seems to work:

result = SumAllClassTwo(input_df).sum_all_class_two()
Related