<commit_before># coding :utf-8

"""
定义一些可以扩展的数据结构

方便序列化/相互转换

"""

import itertools
from functools import reduce

import numpy as np
import pandas as pd
import six
import talib
from QUANTAXIS.QAData.data_fq import QA_data_stock_to_fq
from QUANTAXIS.QAData.data_resample import QA_data_tick_resample
from QUANTAXIS.QAIndicator import EMA, HHV, LLV, SMA
from QUANTAXIS.QAUtil import (QA_Setting, QA_util_log_info,
                              QA_util_to_json_from_pandas)

# ATR


class __stock_hq_base():
    def __init__(self, DataFrame):
        self.data = DataFrame
        self.type = ''
        self.if_fq = 'bfq'
        self.mongo_coll = QA_Setting.client.quantaxis
        self.open = DataFrame['open']
        self.high = DataFrame['high']
        self.low = DataFrame['low']
        self.close = DataFrame['close']
        if 'volume' in DataFrame.columns:
            self.vol = DataFrame['volume']
        else:
            self.vol = DataFrame['vol']
        if 'date' in self.data.index.names.index:
            self.date = self.data.index.levels[self.data.index.names.index(
                'date')]
            self.datetime = self.date
        elif 'datetime' in self.data.index.names.index:
            self.datetime = self.data.index.levels[self.data.index.names.index(
                'datetime')]
            self.date = self.datetime.apply(lambda x: str(x)[0:10])
        self.index = DataFrame.index
        self.code = self.data.index.levels[self.data.index.names.index('code')]

    def len(self):
        return len(self.data)

    def reverse(self):
        return __stock_hq_base(self.data[::-1])

    def show(self):
        return QA_util_log_info(self.data)

    def query(self, query_text):
        return self.data.query(query_text)

    def to_list(self):
        return np.asarray(self.data).tolist()

    def to_pd(self):
        return self.data

    def to_numpy(self):
        return np.asarray(self.data)

    def to_json(self):
        return QA_util_to_json_from_pandas(self.data)

    def sync_status(self, __stock_hq_base):
        '固定的状态要同步 尤其在创建新的datastruct时候'
        (__stock_hq_base.if_fq, __stock_hq_base.type, __stock_hq_base.mongo_coll) = (
            self.if_fq, self.type, self.mongo_coll)
        return __stock_hq_base

    def splits(self):
        if self.type in ['stock_day', 'index_day']:
            return list(map(lambda data: self.sync_status(data), list(map(lambda x: (
                self.data[self.data['code'] == x].set_index(['date', 'code'], drop=False)), self.code))))
        elif self.type in ['stock_min']:
            return list(map(lambda data: self.sync_status(data), list(map(lambda x: (
                self.data[self.data['code'] == x].set_index(['datetime', 'code'], drop=False)), self.code))))

    def add_func(self, func, *arg, **kwargs):
        return self.sync_status(__stock_hq_base(pd.concat(list(map(lambda x: func(
            self.data[self.data['code'] == x], *arg, **kwargs), self.code)))))
    def pivot(self,column_):
        try:
            return self.data.pivot_table(index='datetime',columns='code',values=column_)
        except:
            return self.data.pivot_table(index='date',columns='code',values=column_)

class QA_DataStruct_Index_day(__stock_hq_base):
    '自定义的日线数据结构'

    def __init__(self, DataFrame):
        self.data = DataFrame
        self.type = 'index_day'
        self.if_fq = ''
        self.mongo_coll = QA_Setting.client.quantaxis.stock_day
        self.open = DataFrame['open']
        self.high = DataFrame['high']
        self.low = DataFrame['low']
        self.close = DataFrame['close']
        if 'volume' in DataFrame.columns:
            self.vol = DataFrame['volume']
        else:
            self.vol = DataFrame['vol']
        self.date = self.data.index.levels[self.data.index.names.index('date')]<commit_msg>pivot with only str<commit_after># coding :utf-8

"""
定义一些可以扩展的数据结构

方便序列化/相互转换

"""

import itertools
from functools import reduce

import numpy as np
import pandas as pd
import six
import talib
from QUANTAXIS.QAData.data_fq import QA_data_stock_to_fq
from QUANTAXIS.QAData.data_resample import QA_data_tick_resample
from QUANTAXIS.QAIndicator import EMA, HHV, LLV, SMA
from QUANTAXIS.QAUtil import (QA_Setting, QA_util_log_info,
                              QA_util_to_json_from_pandas)

# ATR


class __stock_hq_base():
    def __init__(self, DataFrame):
        self.data = DataFrame
        self.type = ''
        self.if_fq = 'bfq'
        self.mongo_coll = QA_Setting.client.quantaxis
        self.open = DataFrame['open']
        self.high = DataFrame['high']
        self.low = DataFrame['low']
        self.close = DataFrame['close']
        if 'volume' in DataFrame.columns:
            self.vol = DataFrame['volume']
        else:
            self.vol = DataFrame['vol']
        if 'date' in self.data.index.names.index:
            self.date = self.data.index.levels[self.data.index.names.index(
                'date')]
            self.datetime = self.date
        elif 'datetime' in self.data.index.names.index:
            self.datetime = self.data.index.levels[self.data.index.names.index(
                'datetime')]
            self.date = self.datetime.apply(lambda x: str(x)[0:10])
        self.index = DataFrame.index
        self.code = self.data.index.levels[self.data.index.names.index('code')]

    def len(self):
        return len(self.data)

    def reverse(self):
        return __stock_hq_base(self.data[::-1])

    def show(self):
        return QA_util_log_info(self.data)

    def query(self, query_text):
        return self.data.query(query_text)

    def to_list(self):
        return np.asarray(self.data).tolist()

    def to_pd(self):
        return self.data

    def to_numpy(self):
        return np.asarray(self.data)

    def to_json(self):
        return QA_util_to_json_from_pandas(self.data)

    def sync_status(self, __stock_hq_base):
        '固定的状态要同步 尤其在创建新的datastruct时候'
        (__stock_hq_base.if_fq, __stock_hq_base.type, __stock_hq_base.mongo_coll) = (
            self.if_fq, self.type, self.mongo_coll)
        return __stock_hq_base

    def splits(self):
        if self.type in ['stock_day', 'index_day']:
            return list(map(lambda data: self.sync_status(data), list(map(lambda x: (
                self.data[self.data['code'] == x].set_index(['date', 'code'], drop=False)), self.code))))
        elif self.type in ['stock_min']:
            return list(map(lambda data: self.sync_status(data), list(map(lambda x: (
                self.data[self.data['code'] == x].set_index(['datetime', 'code'], drop=False)), self.code))))

    def add_func(self, func, *arg, **kwargs):
        return self.sync_status(__stock_hq_base(pd.concat(list(map(lambda x: func(
            self.data[self.data['code'] == x], *arg, **kwargs), self.code)))))
    def pivot(self,column_):
        assert isinstance(column_,str)
        try:
            return self.data.pivot(index='datetime',columns='code',values=column_)
        except:
            return self.data.pivot(index='date',columns='code',values=column_)

class QA_DataStruct_Index_day(__stock_hq_base):
    '自定义的日线数据结构'

    def __init__(self, DataFrame):
        self.data = DataFrame
        self.type = 'index_day'
        self.if_fq = ''
        self.mongo_coll = QA_Setting.client.quantaxis.stock_day
        self.open = DataFrame['open']
        self.high = DataFrame['high']
        self.low = DataFrame['low']
        self.close = DataFrame['close']
        if 'volume' in DataFrame.columns:
            self.vol = DataFrame['volume']
        else:
            self.vol = DataFrame['vol']
        self.date = self.data.index.levels[self.data.index.names.index('date')]
{'fraction_non_alphanumeric': 0.12843681425067768, 'fraction_numerical': 0.001290822253775655, 'mean_word_length': 2.6239476145930776, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5276762', 'n_tokens_mistral': 2681, 'n_tokens_neox': 2638, 'n_words': 466}