file_writer.py

class FileWriter:
    # TODO: cleanup output_file
    def __init__(self, data, out_format=None, country="Macedonia"):
        self.data = data
        self.format = out_format
        self.country = country

    def output_file(self):
        """
        Write list of hotels in file
        :return:
        """
        format = self.format.lower()
        file_name = ""
        if format == "json" or format is None:
            import json

            file_name = "hotels-in-{country}.txt".format(
                country=self.country.replace(" ", "-")
            )
            with open(file_name, "w") as outfile:
                json.dump(list(self.data), outfile, indent=2, ensure_ascii=False)
        elif format == "excel":
            from openpyxl import Workbook

            wb = Workbook()
            ws = wb.active

            heading1 = "#"
            heading2 = "Accommodation"
            ws.cell(row=1, column=1).value = heading1
            ws.cell(row=1, column=2).value = heading2

            for i, item in enumerate(self.data):
                # Extract number and title from string
                tokens = item.split()
                n = tokens[0]
                title = " ".join(tokens[2:])

                ws.cell(row=i + 2, column=1).value = n
                ws.cell(row=i + 2, column=2).value = title

            file_name = "hotels-in-{country}.xls".format(
                country=self.country.replace(" ", "-")
            )
            wb.save(file_name)
        elif format == "csv":
            file_name = "hotels-in-{country}.csv".format(
                country=self.country.replace(" ", "-")
            )
            with open(file_name, "w") as outfile:
                for i, item in enumerate(self.data):
                    # Extract number and title from string
                    tokens = item.split()
                    n = tokens[0]
                    title = " ".join(tokens[2:])

                    s = n + ", " + title + "\n"
                    outfile.write(s)

        return file_name