match_spider/ybw/utils/RedisDB.py

# -*- coding: utf-8 -*-
"""
Created on 2024-02-27
---------
@summary: redis 去重
---------
@author: Lzz
"""
import hashlib

import redis


class RedisFilter:
    redis_db = None

    def __init__(self, redis_url=None, expire_time=None):

        self.__class__.redis_db = redis.StrictRedis.from_url(redis_url)  # 单机

        self._ex = expire_time or 86400 * 365 * 1  # 1年 = 86400 * 365 * 1

    def __repr__(self):
        return "<RedisFilter: {}>".format(self.redis_db)

    def exists(self, key):
        """全量检索"""
        if self.redis_db.exists(key) > 0:
            return True
        return False

    def add(self, keys, *args, **kwargs):
        """
        添加数据  删除数据：redis_db.delete("pylist_" + key)
        @param keys: 检查关键词在 redis 中是否存在，支持列表批量
        @return: list / 单个值(如果数据已存在 返回 False 否则返回 True, 可以理解为是否添加成功)
        """
        is_list = isinstance(keys, list)
        keys = keys if is_list else [keys]

        is_added = []
        for key in keys:
            if not self.exists(key):
                is_added.append(self.redis_db.set(key, 1, ex=self._ex))
            else:
                is_added.append(False)

        return is_added if is_list else is_added[0]

    def get(self, keys):
        """
        检查数据是否存在
        @param keys: list / 单个值
        @return: list / 单个值 （存在返回True 不存在返回False)
        """
        is_list = isinstance(keys, list)
        keys = keys if is_list else [keys]

        is_exist = []
        for key in keys:
            is_exist.append(self.exists(key))

        # 判断数据本身是否重复
        temp_set = set()
        for i, key in enumerate(keys):
            if key in temp_set:
                is_exist[i] = True
            else:
                temp_set.add(key)

        return is_exist if is_list else is_exist[0]


def get_sha256(*args):
    """
    @summary: 获取唯一的64位值， 用于获取唯一的id
    ---------
    @param *args: 参与联合去重的值
    ---------
    @result: 5580c91ea29bf5bd963f4c08dfcacd983566e44ecea1735102bc380576fd6f30
    """

    sha256 = hashlib.sha256()
    for arg in args:
        sha256.update(str(arg).encode())
    return sha256.hexdigest()  # 64位


def rexists(dedup, data):
    data = [data] if not isinstance(data, list) else data
    args = sorted(data)
    pykey = "pylist_" + get_sha256(*args)
    if dedup.get(pykey):
        ''' 存在 '''
        return True
    else:
        ''' 不存在 '''
        return False


def radd(dedup, data):
    data = [data] if not isinstance(data, list) else data
    args = sorted(data)
    pykey = "pylist_" + get_sha256(*args)
    state = dedup.add(pykey)
    return state