modify class to be more python