Parse a CSV file and format the data

Posted on 27-Aug-2022

Here is a script that reads a CSV file, extract some data from it and generates another CSV file with the result

#!/usr/bin/env python3

import json
import csv


# This class will take care only of reading and writing to files.
# Adding multiple methods for different file formats will be easier
class FileManager:

    @staticmethod
    def parse_csv_file(filename):
        # open file for read
        with open(filename, encoding='utf-8') as fh:
            # parse csv content and return the content as a list
            return list(csv.reader(fh, delimiter=',', quotechar='"'))

    @staticmethod
    def write_json_file(filename, content):
        # open file for writing. Notice "w" as second parameter
        with open(filename, 'w', encoding='utf-8') as file:
            # convert the python list to a json string and write it to the file
            file.write(json.dumps(content))

    @staticmethod
    def write_csv_file(filename, content):
        # if content is empty, do not create csv file
        if len(content) < 1:
            return False
        # open file for writing. Notice "w" as second parameter
        with open(filename, 'w', encoding='utf-8') as file:
            # create CSV writer and get the header from the first element keys
            writer = csv.DictWriter(file, fieldnames=content[0].keys())
            # write all content to the csv file
            writer.writerows(content)


class DataTransformer:
    def __init__(self):
        # save csv header. Not really necessary for this example
        self.header = []
        # save all required_data that will be the output of the script logic.
        # you can replace this with a local list in the method and return directly
        self.required_data = []

    @staticmethod
    def format_output_item(csv_line):
        # extract only the data that you want from the csv line
        return {
            'name': csv_line[0],
            'birth_year': csv_line[6],
            'gender': csv_line[8],
            'homeworld': csv_line[9],
        }

    def parse(self, filename):
        # parse the csv file and get the data in a list format
        csv_lines = FileManager.parse_csv_file(filename)
        # iterate over each csv line
        for index, line in enumerate(csv_lines):
            # print(index, item)
            if index == 0:
                self.header = line
                continue
            # create the output structure as you want
            formatted_item = self.format_output_item(line)
            # append the resulted item to the required data
            self.required_data.append(formatted_item)


if __name__ == '__main__':
    dt = DataTransformer()
    dt.parse('starwars.csv')
    # print(dt.required_data)
    FileManager.write_csv_file('export.csv', dt.required_data)