【python爬虫】设计自己的爬虫 3. 文件数据保存封装

考虑到爬取的多媒体文件要保存到本地，因此封装了一个类来专门处理这样的问题，下面看代码：

class FileStore:
    def __init__(self, file_path, read_file_mode='r',write_file_mode='wb'):
        """
        初始化 FileStore 实例

        Parameters:
        - file_path (str): 文件路径
        - read_file_mode (str): 读取文件时的模式，默认为 'r'
        - write_file_mode (str): 写入文件时的模式，默认为 'wb'
        """
        self.file_path = file_path
        self.read_file_mode = read_file_mode
        self.write_file_mode = write_file_mode
        self.file_type = os.path.splitext(file_path)[1][1:]

    def read(self):
        """
        读取文件内容

        Returns:
        - result: 读取到的文件内容，根据文件类型可能是列表、字典等
        """
        if self.file_type == CSV_TYPE:
            df = pd.read_csv(self.file_path)
            # 调用df的values属性，再调用tolist方法，可将数据转化为列表形式
            result = df.values.tolist()
            # 也可对df进行逐行遍历，得到列表类型的结果
            # for index, row in df.iterrows():
            #     logging.info(row.tolist())

        elif self.file_type == JSON_TYPE:
            # ensure_ascii 中文编码
            result = json.load(open(self.file_path, encoding='utf-8'))
        else:
            with open(self.file_path, mode=self.read_file_mode, encoding='utf-8') as file:
                result = file.read()

        return result

    # 还要加入对file_data类型的判断
    # 一般默认file_data类型为dict,后面有不同的类型再完善
    def write(self, file_data):
        """
        写入文件内容

        Parameters:
        - file_data: 要写入的数据，根据文件类型可能是字典、列表、字节流等
        """
        if self.file_type == CSV_TYPE and isinstance(file_data, dict):
            df = pd.DataFrame(file_data, columns=file_data.keys(), index=file_data.keys())
            df.to_csv(self.file_path, index=False)
        elif self.file_type == CSV_TYPE and isinstance(file_data, list):
            with open(self.file_path, self.write_file_mode, newline='', encoding='utf-8') as csvfile:
                fieldnames = file_data[0].keys()
                csvwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
                # 写入列名
                csvwriter.writeheader()
                # 写入数据
                csvwriter.writerows(file_data)

        elif self.file_type == JSON_TYPE and isinstance(file_data, dict):
            # ensure_ascii 中文编码
            json.dump(file_data, open(self.file_path, self.write_file_mode, encoding='utf-8'), indent=2, ensure_ascii=False)
        elif self.file_type == TXT_TYPE and isinstance(file_data, dict):
            with open(self.file_path, self.write_file_mode, encoding='utf-8') as file:
                for item_key, item_value in file_data.items():
                    file.write(f'{item_key}:{item_value}\n')
        # 如果传入的是字节流
        elif isinstance(file_data, bytes):
            with open(self.file_path, self.write_file_mode) as file:
                file.write(file_data)
        else:
            # 写入流媒体 图片、视频等
            with open(self.file_path, self.write_file_mode) as file:
                # 使用流式下载
                for chunk in file_data.iter_content(chunk_size=1024):
                    if chunk:
                        file.write(chunk)

使用的示例代码如下

if __name__ == "__main__":
    file = FileStore('../../files/data.txt')
    data = {
        'name': 'joe',
        'key': 'great'
    }
    file.write(data)
    # file.read()

文件类型的枚举如下

CSV_TYPE = 'csv'
JSON_TYPE = 'json'
TXT_TYPE = 'txt'

要注意的是设置初始化的时候要读取文件和写入文件的模式，如果模式不匹配会读取或写入的时候会报错