taiqiuxuanzhi_spider/store.py

import os
import pandas as pd
from log import PPLogger
from setting import CSV_FILENAME, CSV_DIR
from threading import Lock
pd.options.display.max_colwidth = 100000


class WuBaStore:
    def __init__(self, method):
        """
        持久化
        :param method: 持久化方式 :'csv',,,
        """
        self.method = method
        self.logger = PPLogger(name='store')
        self.logger.setup_logger()
        self.lock = Lock()

    def run(self, *args):
        if self.method == 'csv':
            self.to_csv(args[0])
        else:
            pass

    def to_csv(self, data_dict: dict, file_dir: str = CSV_DIR, file_name: str = CSV_FILENAME):
        if not os.path.exists(file_dir):
            os.mkdir(file_dir)
        try:
            file_path = os.path.join(file_dir, file_name)
            data = pd.DataFrame([data_dict])
            with self.lock:
                if not os.path.exists(file_path):
                    data.to_csv(path_or_buf=file_path, index=False, header=True, encoding='utf8')
                else:
                    data.to_csv(path_or_buf=file_path, index=False, header=False, encoding='utf8', mode='a')
        except Exception as e:
            self.logger.error(e)