魔のkyo的BLOG

import re


def bind(func, *args, **kw):
    return lambda *_args, **_kw: func(*args, *_args, **kw, **_kw)


class Token:
    def __init__(self, type, value):
        self.type = type
        self.value = value


def tokenize(expression):
    # 词法规则
    token_specification = [
        ('STOCKSET', r'[0-9a-zA-Z]+\.[0-9a-zA-Z]+'),  # 股票集，指数代码或自定义股票组代码
        ('PLUS', r'\+'),  # 匹配加号
        ('MINUS', r'-'),  # 匹配减号
        ('INTERSECT', r'\&'),  # 匹配&
        ('LPAREN', r'\('),  # 匹配左括号
        ('RPAREN', r'\)'),  # 匹配右括号
        ('WS', r'\s+'),  # 匹配空格，忽略
    ]
    tok_regex = '|'.join('(?P<%s>%s)' % pair for pair in token_specification)
    token_list = []
    for mo in re.finditer(tok_regex, expression):
        kind = mo.lastgroup
        value = mo.group()
        if kind == 'WS':
            continue  # 忽略空格
        token_list.append(Token(kind, value))
    return token_list


class ASTNode:
    def __init__(self, type, value=None, left=None, right=None):
        self.type = type
        self.value = value
        self.left = left
        self.right = right


def parse_expression(tokens):
    def parse_term():
        if tokens[0].type == 'STOCKSET':
            value = tokens.pop(0).value
            return ASTNode('STOCKSET', value=value)
        elif tokens[0].type == 'LPAREN':
            tokens.pop(0)  # 去掉左括号
            node = parse_expression(tokens)
            if tokens[0].type != 'RPAREN':
                raise SyntaxError("Invalid syntax")
            tokens.pop(0)  # 去掉右括号
            return node
        else:
            raise SyntaxError("Invalid syntax")

    node = parse_term()
    while tokens:
        if tokens[0].type == 'RPAREN':
            break
        if tokens[0].type not in ('PLUS', 'MINUS', 'INTERSECT'):
            raise SyntaxError("Invalid syntax")
        op = tokens.pop(0).type
        right = parse_term()
        node = ASTNode(op, left=node, right=right)
    return node


def evaluate_ast(node, code2set):
    if node.type == 'STOCKSET':
        return code2set(node.value)
    elif node.type == 'PLUS':
        return evaluate_ast(node.left, code2set).union(evaluate_ast(node.right, code2set))
    elif node.type == 'MINUS':
        return evaluate_ast(node.left, code2set) - evaluate_ast(node.right, code2set)
    elif node.type == 'INTERSECT':
        return evaluate_ast(node.left, code2set) & evaluate_ast(node.right, code2set)


class StockSetExpression:
    # exp_str: 表达式的字符串
    def __init__(self, exp_str):
        exp_str = exp_str
        tokens = tokenize(exp_str)  # 词法分析，得到token序列
        self.ast = parse_expression(tokens)   # 语法分析，得到抽象语法树

    # 表达式求值
    # code2set: 指数/股票组代码到股票代码集合的转换方法
    def evaluate(self, code2set):
        return evaluate_ast(self.ast, code2set)     # 遍历抽象语法树计算表达式的值


# test
if __name__ == "__main__":
    

    expression = "(000300.SH + 000905.SH) & 000852.SH - 399006.SZ"

    # 得到股票成分股集合
    def get_index_con(index_code, trade_date):
        if index_code == "000300.SH":
            return {"000001.SZ", "000002.SZ", "000003.SZ"}
        if index_code == "000905.SH":
            return {"000002.SZ", "000003.SZ", "000004.SZ"}
        if index_code == "000852.SH":
            return {"000003.SZ", "000004.SZ", "000005.SZ"}
        if index_code == "399006.SZ":
            return {"000004.SZ", "000005.SZ", "000006.SZ"}

    # 创建一个表达式对象，实际上就是做好词法分析语法分析，把抽象语法树结构建出来
    exp_obj = StockSetExpression(expression)
    print("Result of {}: {}".format(
        expression,
        # 表达式求值，实际上就是遍历树，传入的是叶子结点的求值方式，也就是具体的指数如何取得当日成分股
        exp_obj.evaluate(bind(get_index_con, trade_date='2024-12-17'))
        ))

    # 这样就把建立树结构和每天求表达式的值分离，从而不需要每天都重新parse表达式
    # 同时又把具体的求指数成分股的函数从表达式求值的过程中分离成可以从外部指定

import colorsys

def generate_colors(n):
    """
    生成 n 个尽量有明显区别的颜色。

    参数：
    n (int)：要生成的颜色数量，必须是正整数。

    返回：
    list：包含 n 个 RGB 颜色值的列表，每个颜色值是一个三元组 (r, g, b)。
    """
    colors = []
    cnt = 0
    k = (n + 9) // 10
    for i in range(k):
        for j in range(10):
            h = (360 / 10) * j + 23 * i
            s = 1.0 if cnt < 25 else 0.8
            l = 0.3 + (j * 19) % 7 / 7 * 0.4
            cnt += 1
            r, g, b = colorsys.hls_to_rgb(h / 360, l, s)
            colors.append((int(r * 255), int(g * 255), int(b * 255)))
            if cnt == n:
                break
    return colors

import matplotlib.pyplot as plt
import numpy as np

colors = generate_colors(30)

fig, ax = plt.subplots(figsize=(6, 6))
for i, color in enumerate(colors):
    rect = plt.Rectangle((0, -i), 1, 1, color=np.array(color) / 255)
    ax.add_patch(rect)
ax.set_xlim(0, 1)
ax.set_ylim(-len(colors), 0)
ax.axis('off')
plt.show()

import time
import threading
from collections import deque
from functools import wraps
from contextlib import contextmanager


# 使用孩子兄弟树存储调用关系（对于递归的情况特判）
class Node:
    def __init__(self, name, parent):
        self.name = name
        self.parent = parent
        self.sibling = None
        self.child = None
        self.total_calls = 0 # 调用计数
        self.recur_counter = 0 # 递归层数
        self.start_ts = 0
        self.total_ts = 0
        self.is_recursive = False

    def get_sub_node(self, name):
        p = self.child
        if p is None:
            q = Node(name, self)
            self.child = q
            return q
        else:
            # 找到name节点并返回，如果找不到，退出循环时p是当前节点的最后一个孩子
            while True:
                if p.name == name:
                    return p
                if p.sibling is None:
                    break
                p = p.sibling
            # 创建新节点
            q = Node(name, self)
            # 把新节点作为p的兄弟，即新的最后一个孩子
            p.sibling = q
            return q

    def call(self):
        self.total_calls += 1
        self.recur_counter += 1
        if self.recur_counter == 1:
            self.start_ts = time.perf_counter()

    def returns(self):
        if self.recur_counter == 1:
            self.total_ts += time.perf_counter() - self.start_ts
        self.recur_counter -= 1

    # 返回名字=name的祖先节点或当前节点自身
    # 如果没有名字=name的祖先并且自身名字也不是name，则返回None
    # 此方法用来辅助处理递归调用
    def ancestor_or_self(self, name):
        p = self
        while p:
            if p.name == name:
                return p
            p = p.parent
        return None


class Profiler:
    def __init__(self):
        self.node_stack = deque()
        self.roots = {}
        self.current_node = None
        self.main_tid = threading.get_ident()

    def begin_range(self, name):
        # 每个线程需要创建一个root节点
        tid = threading.get_ident()
        if tid not in self.roots:
            root_name = f"thread:{tid}" if tid != self.main_tid else "main_thread"
            self.current_node = self.roots[tid] = Node(root_name, None)

        self.node_stack.append(self.current_node)
        p = self.current_node.ancestor_or_self(name)
        if p: # 说明发生了递归
            self.current_node = p
        else:
            self.current_node = self.current_node.get_sub_node(name)
        self.current_node.call()

    def end_range(self):
        self.current_node.returns()
        self.current_node = self.node_stack.pop()


class ProfilerViewer:
    def __init__(self, root, view_func):
        self.root = root
        self.view_func = view_func

    def view(self):
        self._view(self.root, 0)

    def _view(self, p, depth):
        self.view_func(p, depth)
        if p.child:
            self._view(p.child, depth+1)
        if p.sibling:
            self._view(p.sibling, depth)


def print_profile(profiler, print_func=print):
    def view_func(p, depth):
        if p.parent and p.parent.total_ts:
            pct = p.total_ts / p.parent.total_ts
        else:
            pct = 1.0
        s = r"{}{} {:.2f}pct {:.3f}ms {}calls".format(' '*depth, p.name, pct*100, p.total_ts*1000, p.total_calls)
        print_func(s)

    for root in profiler.roots.values():
        viewer = ProfilerViewer(root, view_func)
        viewer.view()


# 通过装饰器的方式统计一个函数的耗时和调用次数
def profile(profiler):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            profiler.begin_range(func.__name__)
            ret = func(*args, **kwargs)
            profiler.end_range()
            return ret
        return wrapper
    return decorator


# 通过with的方式统计一个代码块的耗时和调用次数
@contextmanager
def profile_scope(profiler, name):
    try:
        profiler.begin_range(name)
        yield
    except Exception:
        raise
    finally:
        profiler.end_range()


profiler = Profiler()


if __name__ == '__main__':

    @profile(profiler)
    def fib(n):
        if n == 1 or n == 2:
            return 1

        with profile_scope(profiler, "sleep"):
            time.sleep(0.1) # 模拟函数实现中一处耗时很多的操作

        # 递归调用自身
        return fib(n-1) + fib(n-2)

    @profile(profiler)
    def add(a, b):
        return a + b

    th = threading.Thread(target=add, args=(1, 2))
    th.start()

    @profile(profiler)
    def test():
        # 计算 fib(5) 
        print(fib(5))

        # 计算10000次add
        for i in range(10000):
            add(1,i)

    test()

    print_profile(profiler)

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)