魔のkyo的BLOG

import redis
from dataclasses import dataclass

@dataclass
class Message: 
    key: str = ""
    value: bytes = b""

class RedisMQProducer:
    def __init__(self, *args, **kw):
        self.client = redis.Redis(*args, **kw)

    def send(self, topic: str, key: str, value: bytes):
        redis_key = f"MQ_{topic}"
        if len(key) >= 32:
            raise Exception("key超长")
        message = key.encode() + (b'\0' * (32 - len(key))) + value
        self.client.lpush(redis_key, message)


def create_redis_mq_producer(server):
    if server.find(':') != -1:
        host, port = server.split(":", 1)
    else:
        host = server
        port = 6379
    return RedisMQProducer(host=host, port=port, db=0)


class RedisMQConsumer:
    def __init__(self, *args, **kw):
        self.client = redis.Redis(*args, **kw)
        self.topics = []


    def subscribe(self, topics):
        self.topics = topics

    def __iter__(self):
        self.idx = 0
        return self

    # 阻塞地获取下一个新消息
    def __next__(self):
        if not self.topics:
            raise StopIteration
        while True:
            topic = self.topics[self.idx]
            self.idx = (self.idx + 1) % len(self.topics)
            redis_key = f"MQ_{topic}"
            data = self.client.rpop(redis_key)
            if data:
                return Message(key=data[:32].decode().strip('\0'), value=data[32:])

    # 获取新消息并返回，直到所有订阅的topic中都没有新消息或者获取到了数量max_records个消息
    def get_messages(self, max_records=20):
        ret = []
        topics = self.topics
        while topics:
            topics_bak = topics
            topics = []
            for topic in topics_bak:
                redis_key = f"MQ_{topic}"
                data = self.client.rpop(redis_key)
                if data:
                    try:
                        ret.append(Message(key=data[:32].decode().strip('\0'), value=data[32:]))
                    except Exception as e:
                        pass
                    topics.append(topic)
                    if len(ret) >= max_records:
                        return ret
        return ret


def create_redis_mq_consumer(server):
    if server.find(':') != -1:
        host, port = server.split(":", 1)
    else:
        host = server
        port = 6379
    return RedisMQConsumer(host=host, port=port, db=0)

import kafka  # pip install kafka-python
import time
import uuid
import datetime
import queue
from dataclasses import dataclass

@dataclass
class Message: 
    key: str = ""
    value: bytes = b""

class KafkaProducer(kafka.KafkaProducer):
    def __init__(self, *args, **kw):
        super().__init__(*args, **kw)

    def send(self, topic: str, key: str, value: bytes):
        # 生成一个长32位的id用于去重，同时顺便记录消息生成时间
        # 默认情况下Kafka保证消息不丢失，但并不能保证不重复，所以我们通过id来自行去重
        id = "{}{}".format(datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3], uuid.uuid4().hex[:15])
        super().send(topic, key=key.encode(), value=id.encode()+value)


def create_kafka_producer(server):
    retry_times = 3
    for i in range(1, retry_times+1):
        try:
            producer = KafkaProducer(
                bootstrap_servers=[server],
                api_version=(2,3,1),
                api_version_auto_timeout_ms=5000
                )
            return producer
        except Exception as e:
            if i == retry_times:
                raise e
            else:
                time.sleep(0.1)


class KafkaConsumer(kafka.KafkaConsumer):
    def __init__(self, **kw):
        super().__init__(**kw)
        self.q = queue.Queue()
        self.s = set()
        self.DUP_DETECT_SIZE = 1000 # 去重检测大小

    # 从最新消息开始订阅（subscribe在父类有实现）
    # def subscribe(self, topic):
    #     pass

    # 实现从一个历史时间点进行消息订阅（能订阅到的消息取决于Kafka服务器配置的保留策略，基于目前的配置可以保证72小时内的消息可以重复消费）
    def subscribe_from_datetime(self, partition, topic, dt):
        if type(dt) is int or type(dt) is float:
            ts = dt
        elif isinstance(dt, datetime.datetime):
            ts = dt.timestamp()
        else:
            ts = time.mktime(time.strptime(f"{dt}", r"%Y-%m-%d %H:%M:%S"))
        offset = self._get_offset_for_time(partition, topic, ts)
        tp = kafka.TopicPartition(topic, partition)
        super().assign([tp])
        super().seek(tp, offset)

    def __iter__(self):
        return self

    # 重新封装阻塞取消息方式，增加去重
    def __next__(self):
        while True:
            message = super().__next__()
            key = message.key
            if not key:
                continue
            id = message.value[:32]
            if id in self.s: # 重复的消息
                continue
            if len(self.s) >= self.DUP_DETECT_SIZE:
                e = self.q.get()
                self.s.remove(e)
            self.s.add(id)
            self.q.put(id)

            return Message(key=key.decode(), value=message.value[32:])

    # 重新封装非阻塞取消息方式，增加去重
    def get_messages(self, max_records=20):
        r = super().poll(timeout_ms=max_records*25, max_records=max_records)
        ret = []
        for messages in r.values():
            for message in messages:
                key = message.key
                if not key:
                    continue
                id = message.value[:32]
                if id in self.s: # 重复的消息
                    continue
                if len(self.s) >= self.DUP_DETECT_SIZE:
                    e = self.q.get()
                    self.s.remove(e)
                self.s.add(uuid)
                self.q.put(uuid)

                ret.append(Message(key=key.decode(), value=message.value[32:]))
        return ret

    def _get_latest_offset(self, partition, topic):
        tp = kafka.TopicPartition(topic, partition)
        super(KafkaConsumer, self).assign([tp])
        off_set_dict = super(KafkaConsumer, self).end_offsets([tp])
        return list(off_set_dict.values())[0]

    def _get_offset_for_time(self, partition, topic, ts):
        tp = kafka.TopicPartition(topic, partition)
        super(KafkaConsumer, self).assign([tp])
        offset_dict = super(KafkaConsumer, self).offsets_for_times({tp: int(ts*1000)})
        offset = list(offset_dict.values())[0]
        if offset is None:
            return self._get_latest_offset(partition, topic)
        else:
            return offset.offset


def create_kafka_consumer(server, group_id):
    retry_times = 3
    for i in range(1, retry_times+1):
        try:
            consumer = KafkaConsumer(
                bootstrap_servers=[server],
                group_id=group_id,
                auto_offset_reset='latest', # earliest, latest
                enable_auto_commit=True,
                api_version=(2,3,1),
                api_version_auto_timeout_ms=5000
                )
            return consumer
        except Exception as e:
            if i == retry_times:
                raise e
            else:
                time.sleep(0.1)

import os
import functools
import redis
from redis.connection import ConnectionPool
import msgpack
import msgpack_numpy
import numpy as np
import pandas as pd
import io

class _HashedSeq(list):
    __slots__ = 'hashvalue'

    def __init__(self, tup, hash=hash):
        self[:] = tup
        self.hashvalue = hash(tup)

    def __hash__(self):
        return self.hashvalue


def _make_key(args, kw, typed=False):
    key = args
    if kw:
        kwd_mark = (object(),)
        key += kwd_mark
        for item in kw.items():
            key += item
    if typed:
        key += tuple(type(v) for v in args)
        if kw:
            key += tuple(type(v) for v in kw.values())
    elif len(key) == 1 and type(key[0]) in {int, str}:
        return key[0]
    return _HashedSeq(key)


def dumps_to_feather(df):
    columns = df.columns
    df.columns = [str(e) for e in df.columns]
    buffer = io.BytesIO()  # 创建一个内存中的字节流缓冲区
    df.to_feather(buffer)  # 将 DataFrame 序列化为 Feather 格式并存储到缓冲区
    buffer.seek(0)  # 重新定位到字节流的开头
    serialized_data = buffer.getvalue()  # 获取序列化后的二进制数据
    return {"data": serialized_data, "columns": list(columns)}


def loads_from_feather(data):
    buffer = io.BytesIO(data["data"])  # 创建一个内存中的字节流缓冲区，并加载序列化的数据
    df = pd.read_feather(buffer)  # 从 Feather 格式的序列化数据中加载 DataFrame
    df.columns = data["columns"]
    return df


def _custom_encode(obj):
    if isinstance(obj, set):
        return {b'__set__': list(obj)}  # 将 set 对象转换为字典形式
    if isinstance(obj, pd.DataFrame):
        return {b'__feather__': dumps_to_feather(obj)}
    return msgpack_numpy.encode(obj)


def _custom_decode(obj):
    if b'__set__' in obj:  # 判断是否是我们之前转换过的 set 对象
        return set(obj[b'__set__'])  # 将字典形式的 set 转换回原始的 set 对象
    if b'__feather__' in obj:
        return loads_from_feather(obj[b'__feather__'])
    return msgpack_numpy.decode(obj)  # 不需要转换的情况下直接返回原始对象


class RedisTTLCache:
    def __init__(self, host='localhost', port=6379):
        self.pool = ConnectionPool(host=host, port=port)

    def cache(self, ttl, scope=os.path.basename(__file__)):
        def decorator(func):
            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                # 构建唯一的缓存键
                cache_key = f"FUNC_CACHE_{(scope, func.__name__, _make_key(args, kwargs))}"
                
                # 尝试从Redis中获取缓存
                redis_client = redis.Redis(connection_pool=self.pool)
                result = redis_client.get(cache_key)
                if result is not None:
                    # 如果缓存存在，直接返回结果
                    result = msgpack.loads(result, object_hook=_custom_decode)
                    return result
                
                # 如果缓存中不存在该结果，则重新计算函数结果
                result = func(*args, **kwargs)
                
                # 将计算结果存入Redis缓存
                redis_client.set(cache_key, msgpack.dumps(result, default=_custom_encode), ex=ttl)
                
                return result
            return wrapper
        return decorator


if __name__ == '__main__':
    redis_ttl_cache = RedisTTLCache(host="localhost")
    ttl_cached = redis_ttl_cache.cache

    @ttl_cached(ttl=5)
    def func(a, b):
        print("Hello world")
        return {
            'result': a + b,
            'colors': {'red', 'green', 'blue'},
            'arr': np.array([[0, 1], [2,3]]),
            'df': pd.DataFrame([[0, 1], [2,3]]),
        }

    print(func(1, 2))
    print(func(1, 2))

# -*- coding: utf-8 -*-
import os
import openpyxl # 解析xlsx
import xlrd # 解析xls
import csv  # 解析csv


"""
类功能概述:
    本代码实现了一个支持读取多种格式（xlsx、xls、csv）Excel 文件的功能。
    主要包含 `Workbook` 和 `Sheet` 相关类，提供了统一的接口来操作不同格式文件。

使用步骤:
    1. 创建 `Workbook` 对象，传入文件路径，根据文件后缀自动选择合适的解析器。
    2. 通过 `Workbook` 对象的 `sheet_by_index` 方法获取指定索引的工作表（`Sheet` 对象）。
    3. 使用 `Sheet` 对象的 `cell_value` 方法获取指定单元格的值，`max_row` 方法获取最大行数，`max_col` 方法获取最大列数。
"""


class Sheet:
    """
    `Sheet` 类是所有工作表类的基类，定义了获取单元格值、最大行数和最大列数的接口。

    方法:
        cell_value(row, col): 获取指定行和列的单元格值。
        max_row(): 获取工作表的最大行数。
        max_col(): 获取工作表的最大列数。
    """

    def cell_value(self, row, col):
        pass

    def max_row(self):
        pass

    def max_col(self):
        pass


class XlsxSheet(Sheet):
    def __init__(self, sheet):
        self.sheet = sheet

    def cell_value(self, row, col):
        return self.sheet.cell(row=row+1, column=col+1).value

    def max_row(self):
        return self.sheet.max_row

    def max_col(self):
        return self.sheet.max_column


class XlsSheet(Sheet):
    def __init__(self, sheet):
        self.sheet = sheet

    def cell_value(self, row, col):
        return self.sheet.cell_value(row, col)

    def max_row(self):
        return self.sheet.nrows

    def max_col(self):
        return self.sheet.ncols


class CsvSheet(Sheet):
    def __init__(self, rows):
        self.rows = rows

    def cell_value(self, row, col):
        if row < len(self.rows) and col < len(self.rows[row]):
            return self.rows[row][col]
        return None

    def max_row(self):
        return len(self.rows)

    def max_col(self):
        max_cols = 0
        for row in self.rows:
            max_cols = max(max_cols, len(row))
        return max_cols


class _Workbook:
    def sheet_by_index(self, index):
        pass


class XlsxFile(_Workbook):
    def __init__(self, path):
        self.workbook = openpyxl.load_workbook(path)

    def sheet_by_index(self, index):
        sheet = self.workbook[self.workbook.sheetnames[index]]
        return XlsxSheet(sheet)


class XlsFile(_Workbook):
    def __init__(self, path):
        self.workbook = xlrd.open_workbook(path)

    def sheet_by_index(self, index):
        sheet = self.workbook.sheet_by_index(index)
        return XlsSheet(sheet)


class CsvFile(_Workbook):
    def __init__(self, path):
        encodings = ['utf-8-sig', 'utf-8']
        for encoding in encodings:
            try:
                with open(path, 'r', encoding=encoding) as file:
                    reader = csv.reader(file)
                    self.rows = list(reader)
                break
            except UnicodeDecodeError:
                continue
        else:
            raise UnicodeDecodeError("无法使用支持的编码读取文件")

    def sheet_by_index(self, index):
        if index == 0:
            return CsvSheet(self.rows)
        return None


class Workbook:
    """
    `Workbook` 类是对外提供的工作簿类，根据文件后缀自动选择合适的解析器。

    属性:
        impl: 具体的工作簿实现对象（`XlsxFile`、`XlsFile` 或 `CsvFile`）。

    方法:
        __init__(path): 初始化工作簿对象，根据文件后缀选择合适的解析器。
        sheet_by_index(index): 获取指定索引的工作表。
    """

    def __init__(self, path):
        filedir, filename = os.path.split(path)
        _, extname = os.path.splitext(filename)
        if extname == ".xlsx":
            self.impl = XlsxFile(path)
        elif extname == ".xls":
            self.impl = XlsFile(path)
        elif extname == ".csv":
            self.impl = CsvFile(path)
        else:
            raise IOError("未知的文件类型")

    def sheet_by_index(self, index):
        return self.impl.sheet_by_index(index)

# -*- coding: utf-8 -*-
import threading
import time
import inspect
import ctypes
import traceback
import sys, os
from functools import wraps

class KillableThread(threading.Thread):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._return_value = None
        self._exception = None

    def _async_raise(tid, exctype):
        """raises the exception, performs cleanup if needed"""
        tid = ctypes.c_long(tid)
        if not inspect.isclass(exctype):
            exctype = type(exctype)
        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
        if res == 0:
            raise ValueError("invalid thread id")
        elif res != 1:
            # """if it returns a number greater than one, you're in trouble,
            # and you should call it again with exc=NULL to revert the effect"""
            ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
            raise SystemError("PyThreadState_SetAsyncExc failed")

    def kill(self):
        KillableThread._async_raise(self.ident, SystemExit)

    def run(self):
        try:
            self._return_value = self._target(*self._args, **self._kwargs)
        except Exception as e:
            self._exception = e

    def join(self):
        super().join()
        if self._exception is not None:
            raise self._exception
        return self._return_value

    def join(self, timeout):
        super().join(timeout)
        if self._exception is not None:
            raise self._exception
        return self._return_value


def _get_thread(tid):
    for t in threading.enumerate():
        if t.ident == tid:
            return t
    return None


def _get_frame_stack(tid):
    for thread_id, stack in sys._current_frames().items():
        if thread_id == tid:
            return stack
    return None


def _get_formated_frame_stack(tid):
    info = []
    th = _get_thread(tid)
    stack = _get_frame_stack(tid)
    info.append('%s thread_id=%d' % (th.name, tid))
    for filename, lineno, _, line in traceback.extract_stack(stack):
        info.append('    at %s(%s:%d)' % (line, filename[filename.rfind(os.path.sep) + 1:], lineno))
    return '\n'.join(info)


def timeout(seconds):
    """
    Decorator to execute a function with a specified timeout.

    Args:
    - seconds (int): The time limit in seconds for the function to complete.

    Returns:
    - function: The decorated function.

    Raises:
    - TimeoutError: If the function does not complete within the specified time limit.

    Usage:
    @timeout(seconds=10)
    def my_function():
        # Function body
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            th = KillableThread(target=func, args=args, kwargs=kwargs)
            th.daemon = True
            th.start()
            ret = th.join(seconds)
            if th.is_alive():
                formated_frame_stack = _get_formated_frame_stack(th.ident)
                th.kill()
                raise TimeoutError(f"{repr(func)} timeout. Frame stack:\n{formated_frame_stack}")
            return ret
        return wrapper
    return decorator


def retry(retries=1, retry_interval=0):
    """
    Decorator to retry a function a specified number of times with a given interval between retries.

    Args:
    - retries (int): The number of times the function should be retried if it raises an exception. If set to 1, the function will be attempted initially and retried once.
    - retry_interval (int): The time interval in seconds to wait between retries.

    Returns:
    - function: The decorated function.

    Raises:
    - The original exception raised by the function if all retries fail.

    Usage:
    @retry(retries=2, retry_interval=2)
    def my_function():
        # Function body
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(retries+1):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if i < retries:
                        time.sleep(retry_interval)
                    else:
                        raise e
        return wrapper
    return decorator