Reformat for PEP-8 and add imports