add funcat

README.md

@@ -0,0 +1,4 @@
+# Funcat
+提供 Python 公式选股。
+
+使用同花顺、通达信等公式，做技术分析，表达十分简洁。

funcat/__init__.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#
+
+import datetime
+
+from .time_series import PriceSeries
+from .func import MovingAverageSeries, CrossOver, minimum, maximum, every, count, hhv, llv, Ref
+from .context import ExecutionContext, symbol, set_current_stock as S, set_current_date as T, set_data_backend
+from .helper import loop
+
+
+# close: CLOSE, C, c
+for name in ["open", "high", "low", "close", "volume"]:
+    cls = type("{}Series".format(name.capitalize()), (PriceSeries, ), {"name": name})
+    obj = cls(dynamic_update=True)
+    for var in [name[0], name[0].upper(), name.upper()]:
+        globals()[var] = obj
+
+
+ma = MA = MovingAverageSeries
+cross = CROSS = CrossOver
+ref = REF = Ref
+MIN = minimum
+MAX = maximum
+EVERY = every
+COUNT = count
+HHV = hhv
+LLV = llv
+
+
+ExecutionContext(date=int(datetime.date.today().strftime("%Y%m%d")),
+                 stock="000001.XSHG",
+                 data_backend=None)._push()

funcat/context.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#
+
+import datetime
+
+import six
+
+from .utils import get_int_date
+
+
+class ExecutionContext(object):
+    stack = []
+
+    def __init__(self, date=None, stock=None, data_backend=None):
+        self._current_date = date
+        self._stock = stock
+        self._data_backend = data_backend
+
+    def _push(self):
+        self.stack.append(self)
+
+    def _pop(self):
+        popped = self.stack.pop()
+        if popped is not self:
+            raise RuntimeError("Popped wrong context")
+        return self
+
+    def __enter__(self):
+        self._push()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    @classmethod
+    def get_active(cls):
+        return cls.stack[-1]
+
+    @classmethod
+    def set_current_date(cls, date):
+        """set current simulation date
+        :param date: string date, "2016-01-04"
+        """
+        if isinstance(date, six.string_types):
+            date = get_int_date(date)
+        elif isinstance(date, datetime.date):
+            date = int(date.strftime("%Y%m%d"))
+        cls.get_active()._current_date = date
+
+    @classmethod
+    def get_current_date(cls):
+        return cls.get_active()._current_date
+
+    @classmethod
+    def set_current_stock(cls, stock):
+        """set current watching stock
+        :param stock: "000002.XSHE"
+        """
+        cls.get_active()._stock = stock
+
+    @classmethod
+    def get_current_stock(cls):
+        return cls.get_active()._stock
+
+    @classmethod
+    def set_data_backend(cls, data_backend):
+        """set current watching stock
+        :param stock: "000002.XSHE"
+        """
+        cls.get_active()._data_backend = data_backend
+
+    @classmethod
+    def get_data_backend(cls):
+        return cls.get_active()._data_backend
+
+
+def set_data_backend(backend):
+    ExecutionContext.set_data_backend(backend)
+
+
+def set_current_stock(stock):
+    ExecutionContext.set_current_stock(stock)
+
+
+def set_current_date(date):
+    ExecutionContext.set_current_date(date)
+
+
+def symbol(order_book_id):
+    """获取股票代码对应的名字
+    :param order_book_id:
+    :returns:
+    :rtype:
+    """
+    data_backend = ExecutionContext.get_data_backend()
+    return data_backend.symbol(order_book_id)

funcat/data/__init__.py

@@ -0,0 +1,4 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#

funcat/data/backend.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#
+
+
+class DataBackend(object):
+    def get_price(self, order_book_id, start, end):
+        """
+        :param order_book_id: e.g. 000002.XSHE
+        :param start: 20160101
+        :param end: 20160201
+        :returns:
+        :rtype: numpy.rec.array
+        """
+        raise NotImplementedError
+
+    def get_order_book_id_list(self):
+        """获取所有的
+        """
+        raise NotImplementedError
+
+    def get_trading_dates(self, start, end):
+        """获取所有的交易日
+
+        :param start: 20160101
+        :param end: 20160201
+        """
+        raise NotImplementedError
+
+    def get_start_date(self):
+        """获取回溯开始时间
+        """
+        raise NotImplementedError
+
+    def symbol(self, order_book_id):
+        """获取order_book_id对应的名字
+        :param order_book_id str: 股票代码
+        :returns: 名字
+        :rtype: str
+        """
+        raise NotImplementedError

funcat/data/tushare_backend.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#
+
+from .backend import DataBackend
+from ..utils import lru_cache, get_str_date_from_int, get_int_date
+
+
+class TushareDataBackend(DataBackend):
+    def __init__(self, start_date="2011-01-01"):
+        import tushare as ts
+        self.start_date = start_date
+        self.ts = ts
+
+    def convert_code(self, order_book_id):
+        return order_book_id.split(".")[0]
+
+    @lru_cache(maxsize=4096)
+    def get_price(self, order_book_id, start, end):
+        """
+        :param order_book_id: e.g. 000002.XSHE
+        :param start: 20160101
+        :param end: 20160201
+        :returns:
+        :rtype: numpy.rec.array
+        """
+        start = get_str_date_from_int(start)
+        end = get_str_date_from_int(end)
+        code = self.convert_code(order_book_id)
+        is_index = False
+        if ((order_book_id.startswith("0") and order_book_id.endswith(".XSHG")) or
+            (order_book_id.startswith("3") and order_book_id.endswith(".XSHE"))
+            ):
+            is_index = True
+        df = self.ts.get_k_data(code, start=start, end=end, index=is_index)
+        df["date"] = df["date"].apply(lambda x: int(x.replace("-", "")))
+        df = df.set_index("date")
+        del df["code"]
+        arr = df.to_records()
+        return arr
+
+    @lru_cache()
+    def get_order_book_id_list(self):
+        """获取所有的股票代码列表
+        """
+        info = self.ts.get_stock_basics()
+        code_list = info.index.sort_values().tolist()
+        order_book_id_list = [
+            (code + ".XSHG" if code.startswith("6") else code + ".XSHE")
+            for code in code_list
+        ]
+        return order_book_id_list
+
+    @lru_cache()
+    def get_trading_dates(self, start, end):
+        """获取所有的交易日
+
+        :param start: 20160101
+        :param end: 20160201
+        """
+        start = get_str_date_from_int(start)
+        end = get_str_date_from_int(end)
+        df = self.ts.get_k_data("000001", index=True, start=start, end=end)
+        trading_dates = [get_int_date(date) for date in df.date.tolist()]
+        return trading_dates
+
+    @lru_cache()
+    def get_start_date(self):
+        """获取回溯开始时间
+        """
+        return self.start_date
+
+    @lru_cache(maxsize=4096)
+    def symbol(self, order_book_id):
+        """获取order_book_id对应的名字
+        :param order_book_id str: 股票代码
+        :returns: 名字
+        :rtype: str
+        """
+        return order_book_id

funcat/func.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# Author: Hua Liang[Stupid ET] <[email protected]>
+#
+
+import numpy as np
+
+from .utils import FormulaException
+from .time_series import PriceSeries, NumericSeries, BoolSeries, fit_series
+
+
+class MovingAverageSeries(NumericSeries):
+    def __init__(self, series, period):
+        import talib
+        if isinstance(series, NumericSeries):
+            series = series.series
+            try:
+                series = talib.MA(series, period)
+            except Exception as e:
+                raise FormulaException(e)
+        super(MovingAverageSeries, self).__init__(series)
+        self.extra_create_kwargs["period"] = period
+
+
+def CrossOver(s1, s2):
+    """s1金叉s2
+    :param s1:
+    :param s2:
+    :returns: bool序列
+    :rtype: BoolSeries
+    """
+    series1, series2 = fit_series(s1.series, s2.series)
+    cond1 = series1 > series2
+    series1, series2 = fit_series(s1[1].series, s2[1].series)
+    cond2 = series1 <= series2  # s1[1].series <= s2[1].series
+    cond1, cond2 = fit_series(cond1, cond2)
+    s = cond1 & cond2
+    return BoolSeries(s)
+
+
+def Ref(s1, n):
+    return s1[n]
+
+
+def minimum(s1, s2):
+    if len(s1) == 0 or len(s2) == 0:
+        raise FormulaException("minimum size == 0")
+    s = np.minimum(s1.series, s2.series)
+    return NumericSeries(s)
+
+
+def maximum(s1, s2):
+    if len(s1) == 0 or len(s2) == 0:
+        raise FormulaException("maximum size == 0")
+    s = np.maximum(s1.series, s2.series)
+    return NumericSeries(s)
+
+
+def count(cond, n):
+    # TODO lazy compute
+    series = cond.series
+    size = len(cond.series) - n
+    try:
+        result = np.full(size, 0, dtype=np.int)
+    except ValueError as e:
+        raise FormulaException(e)
+    for i in range(size - 1, 0, -1):
+        s = series[-n:]
+        result[i] = len(s[s == True])
+        series = series[:-1]
+    return NumericSeries(result)
+
+
+def every(cond, n):
+    return count(cond, n) == n
+
+
+def hhv(s, n):
+    # TODO lazy compute
+    series = s.series
+    size = len(s.series) - n
+    try:
+        result = np.full(size, 0, dtype=np.float64)
+    except ValueError as e:
+        raise FormulaException(e)
+    for i in range(size - 1, 0, -1):
+        s = series[-n:]
+        result[i] = s.max()
+        series = series[:-1]
+    return NumericSeries(result)
+
+
+def llv(s, n):
+    # TODO lazy compute
+    series = s.series
+    size = len(s.series) - n
+    try:
+        result = np.full(size, 0, dtype=np.float64)
+    except ValueError as e:
+        raise FormulaException(e)
+    for i in range(size - 1, 0, -1):
+        s = series[-n:]
+        result[i] = s.min()
+        series = series[:-1]
+    return NumericSeries(result)