【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_pathself.read_file_mode = read_file_modeself.write_file_mode = write_file_modeself.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'

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