魔のkyo的BLOG

import datetime
import uuid

def gen_uuid():
    return "{}{:016x}".format(datetime.datetime.now().strftime("%Y%m%d%H%M%S"), (uuid.uuid4().int&(0xFFFFFFFFFFFFFFFF)) )
    

# -*- coding: utf-8 -*-
import kafka
import time
import datetime
import uuid
from queue import Queue


class KafkaProducer(kafka.KafkaProducer):
    def __init__(self, server):
        super().__init__(
            bootstrap_servers=[server],
            api_version=(2,3,1),
            api_version_auto_timeout_ms=5000
            )

    # topic可以认为是消息的信道，和subscribe时对应
    # key并没有明确的意义，可以根据需要自定义，例如表示消息的类型在后面再根据key进行不同的处理，kafka底层保证相同key的消息在kafka集群时会在相同的分区上处理，从而保证相同key的消息的有序性
    def send(self, topic, key, data):
        assert(isinstance(data, bytes))
        # 用时间+UUID拼成32位唯一值用于去重和在调试时查看发送时间
        timestamp = "{}{}".format(datetime.datetime.now().strftime("%Y%m%d%H%M%S%f")[:-3], uuid.uuid4().hex[:15])
        timestamp = timestamp[:32]
        value = timestamp.encode() + data
        return super().send(topic, key=key.encode(), value=value)


class KafkaConsumer(kafka.KafkaConsumer):
    def __init__(self, server, group_id=None):
        super().__init__(
            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
            )
        self.server = server
        self.q = Queue()
        self.s = set()
        self.DUP_DETECT_SIZE = 1000 # 去重检测大小

    # 从最新消息开始订阅
    def subscribe(self, topic):
        if topic not in super().topics():
            # 创建一个临时的producer发送一个空消息的以便把topic创建出来，key是空的，consumer会丢弃掉
            # 这并不是必要的，我觉得可能是一个BUG，先订阅一个不存在的topic，之后再send消息到这topic上，consumer可能收不到
            # 后续版本可能已经修复了，不清楚是kafka-python的还是kafka的
            producer = kafka.KafkaProducer(
                bootstrap_servers=[self.server],
                api_version=(2,3,1),
                api_version_auto_timeout_ms=5000
            )
            producer.send(topic, key=b"", value=b"")
            producer.flush()
        super().subscribe(topic)

    # 实现从一个历史时间点进行消息订阅（能订阅到的消息取决于Kafka服务器配置的保留策略，基于目前的配置可以保证72小时内的消息可以重复消费）
    def subscribe_from_datetime(self, topic, dt):
        if topic not in super().topics():
            # 创建一个临时的producer发送一个空消息的以便把topic创建出来，key是空的，consumer会丢弃掉
            producer = kafka.KafkaProducer(
                bootstrap_servers=[self.server],
                api_version=(2,3,1),
                api_version_auto_timeout_ms=5000
            )
            producer.send(topic, key=b"", value=b"")
            producer.flush()
        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(topic, ts)
        partition = 0
        tp = kafka.TopicPartition(topic, partition)
        super().assign([tp])
        super().seek(tp, offset)

    def __iter__(self):
        return self

    # 重新封装阻塞取消息方式，增加去重
    # 迭代的元素类型是三元组(key, data, timestamp)
    def __next__(self):
        while True:
            message = super().__next__()
            msg_type = message.key
            if not msg_type or len(message.value) < 32:
                    continue
            timestamp = message.value[:32].decode()
            if timestamp in self.s: # 重复的消息
                continue
            if len(self.s) >= self.DUP_DETECT_SIZE:
                e = self.q.get()
                self.s.remove(e)
            self.s.add(timestamp)
            self.q.put(timestamp)

            data = message.value[32:]
            return (msg_type.decode(), data, timestamp)

    # 重新封装非阻塞取消息方式，增加去重
    # 返回三元组(key, data, timestamp)构成的list
    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:
                msg_type = message.key
                if not msg_type or len(message.value) < 32:
                    continue
                timestamp = message.value[:32].decode()
                if timestamp in self.s: # 重复的消息
                    continue
                if len(self.s) >= self.DUP_DETECT_SIZE:
                    e = self.q.get()
                    self.s.remove(e)
                self.s.add(timestamp)
                self.q.put(timestamp)

                data = message.value[32:]
                ret.append((msg_type.decode(), data, timestamp))
        return ret


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

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

import time
import datetime
from utils.kafka_util import KafkaProducer
import json
import decimal

def to_jsonable(o):
    if o is None:
        return o
    if isinstance(o, int) or isinstance(o, float) or isinstance(o, str) or isinstance(o, bool):
        return o
    if isinstance(o, list) or isinstance(o, tuple) or isinstance(o, set):
        return [to_jsonable(e) for e in o]
    if isinstance(o, dict):
        return {k: to_jsonable(v) for (k,v) in o.items()}
    if isinstance(o, datetime.datetime):
        return o.strftime('%Y-%m-%d %H:%M:%S')
    if isinstance(o, datetime.date):
        return o.strftime('%Y-%m-%d')
    if isinstance(o, datetime.time):
        return o.strftime('%H:%M:%S')
    if isinstance(o, decimal.Decimal):
        return float(o)
    try:
        return { k: to_jsonable(v) for k,v in vars(o).items()}
    except Exception as e:
        return o

producer = KafkaProducer(server="192.168.1.99")

class A:
    def __init__(self):
        self.x = 1
        self.y = 3.14

class B:
    def __init__(self):
        self.x = A()
        self.y = datetime.datetime.now()

producer.send("topic-test", "obj", json.dumps(to_jsonable(B())).encode() )
time.sleep(1)

from utils.kafka_util import KafkaConsumer
import time
import json


consumer = KafkaConsumer(server="192.168.1.99")
consumer.subscribe("topic-test")
while True:
    messages = consumer.get_messages()
    for (key, data, timestamp) in messages:
        print(key, json.loads(data), timestamp)
    time.sleep(0.1)

# -*- coding: utf-8 -*-
import sys
import os
from functools import cmp_to_key


def mycmp_filename(name1, name2):
    if name1.upper() == "README.MD":
        return -1
    elif name2.upper() == "README.MD":
        return 1
    elif name1 < name2:
        return -1
    elif name2 < name1:
        return 1
    return 0


def gen(e, path, level):
    content = ""
    name, ext = os.path.splitext(e)
    if os.path.isfile(path):
        if not name.startswith('_') and ext.lower() == ".md":
            if os.path.getsize(path):
                content += "{}- [{}]({})\n".format("  "*level, name, path.replace('\\', '/') )
    elif os.path.isdir(path):
        if name.startswith('_'):
           return "" 
        content2 = ""
        for e2 in sorted(os.listdir(path), key=cmp_to_key(mycmp_filename)):
            content2 += gen(e2, os.path.join(path, e2), level+1)
        if content2:
            if name:
                content += "{}- {}\n".format("  "*level, name)
            content += content2
        else:
            content = ""
    return content


def main():
    with open("_sidebar.md", "w", encoding="utf-8") as fp:
        fp.write("<!-- docs/_sidebar.md -->\n")
        fp.write("{}\n".format(gen("", ".", -1)))

if __name__ == '__main__':
    main()

import re
import json

# 删除非ascii的编码
def omit_unascii(data):
    new_data = b''
    p = 0
    m = re.search(b"[^\x00-\x7F]+", data[p:])
    while m:
        new_data += data[p:p+m.start()]
        p += m.end()
        m = re.search(b"[^\x00-\x7F]+", data[p:])
    new_data += data[p:]
    return new_data


def json_loads(msg):
    try:
        # 先尝试正常解析（按UTF8解码）
        obj = json.loads(msg)
    except Exception as e:
        try:
            # 尝试先按GBK解码再解析
            obj = json.loads(msg.decode("gb2312", "ignore"))
        except Exception as e:
            try:
                # 尝试删除所有非ascii编码后再解析
                obj = json.loads(omit_unascii(msg))
            except Exception as e:
                raise e
    return obj

# -*- coding: utf-8 -*-
import pymongo
import threading


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)


local = threading.local()

class MongoDB:
    @staticmethod
    def instance(*args, **kw):
        if not hasattr(local, "mongodb") or local.mongodb is None:
            local.mongodb = {}
        key = _make_key(args, kw, typed=False)
        if key not in local.mongodb:
            local.mongodb[key] = MongoDB(*args, **kw)
        return local.mongodb[key]

    def __init__(self, host="127.0.0.1", port=27017, dbname=None, username=None, password=None, **kw):
        self.client = pymongo.MongoClient(f"mongodb://{host}:{port}/{dbname}", username=username, password=password, **kw)
        if dbname:
            self.db = self.client[dbname]
        else:
            self.db = self.client.test

    def close(self):
        return self.client.close()

    def __getattr__(self, key):
        return self.db[key]

    def __getitem__(self, key):
        return self.db[key]

    def has_collection(self, name):
        return name in self.db.list_collection_names()


# 常用运算符
# 逻辑运算
# {$or:[expression1,expression2,...]}
# {$and:[expression1,expression2,...]}
# {$not:expression1}}
# 比较运算
# {field:{$eq:value}}  ==
# {field:{$ne:value}}  !=
# {field:{$lt:value}}  <
# {field:{$lte:value}} <=
# {field:{$gt:value}}  >
# {field:{$gte:value}} >=
# {field:{$in:[value1,value2,...]} in
# {field:{$nin:[value1,value2,...]} not in
# 正则匹配
# {field:{"$regex": "正则表达式"}

# query: 查询条件, ex: {"type": "stock"}
# fields: 提取字段, ex: ["code", "sec_name"]
# sort: 排序字段和排序方式, ex: [("code", -1)] 表示按code排倒序
# limit: 之取前limit个, ex: 10
# **kw: 查询条件，把query展开成参数来写，作用是相同的
def get_securities(query={}, fields=None, sort=None, limit=None, **kw):
    q = gen_query("securities", query, fields, sort, **kw)
    return [doc for doc in (q.limit(limit) if limit else q)]

def gen_query(collection, query={}, fields=None, sort=None, **kw):
    mongo = MongoDB.instance(host="192.168.1.99", dbname="ftresearch", username="ftresearch", password="******")
    q = mongo[collection].find(dict(**query, **kw), to_projection_dict(fields))
    if sort:
        q = q.sort(to_sort_list(sort))
    return q


def to_projection_dict(fields):
    if fields is None:
        return {"_id": 0}
    elif isinstance(fields, dict):
        return fields
    else:
        ret = {"_id": 0}
        for k in fields:
            ret[k] = 1
        return ret


def to_sort_list(sort):
    if not sort:
        return None
    ret = []
    if isinstance(sort, list):
        for e in sort:
            if isinstance(e, str):
                ret.append((e, 1))
            else:
                ret.append(e)
    else:
        if isinstance(sort, str):
            ret.append((sort, 1))
        else:
            ret.append(sort)
    return ret


if __name__ == '__main__':
    # 查询以68编码开头的所有股票，只取code, sec_name字段
    arr = get_securities(code={"$regex": "^68"}, type="stock", fields=["code", "sec_name"])
    for e in arr:
        print(e)

    # 查询所有ST状态和*ST状态的未退市的股票信息，按股票代码排倒序，取前10只
    arr = get_securities(
        query={
        "type": "stock",
        "contract_state": "Active",
        "special_type": {"$in": ["ST", "StarST"]},
        },
        fields=["code", "sec_name", "exchange", "special_type"],
        sort=[("code", -1)],
        limit=10)
    for e in arr:
        print(e)

    # 取2023-01-04日当天在交易状态的股指期货合约代码
    codes = [e["code"]
        for e in get_securities(
            fields=["code"],
            type="future",
            product={"$in": ["IC", "IF", "IH", "IM"]},
            listed_date={"$lte":"2023-01-04"},
            de_listed_date={"$gte":"2023-01-04"}
            )
        ]
    print(codes)