Initial support for polars.Series and polars.DataFrame.

Author: mrjbq7

CHANGELOG

@@ -3,6 +3,12 @@
 
 - [NEW]: Upgrade to Cython 0.29.24
 
+- [NEW]: Support polars.Series in Function API
+
+- [NEW]: Support polars.Series in Streaming API
+
+- [NEW]: Support polars.DataFrame in Abstract API
+
 0.4.21
 ======
 

README.md

@@ -233,7 +233,7 @@ will have an initial "lookback" period (a required number of observations
 before an output is generated) set to ``NaN``.
 
 For convenience, the Function API supports both ``numpy.ndarray`` and
-``pandas.Series`` inputs.
+``pandas.Series`` and ``polars.Series`` inputs.
 
 All of the following examples use the Function API:
 
@@ -270,9 +270,10 @@ If you're already familiar with using the function API, you should feel right
 at home using the Abstract API.
 
 Every function takes a collection of named inputs, either a ``dict`` of
-``numpy.ndarray`` or ``pandas.Series``, or a ``pandas.DataFrame``. If a
-``pandas.DataFrame`` is provided, the output is returned as a
-``pandas.DataFrame`` with named output columns.
+``numpy.ndarray`` or ``pandas.Series`` or ``polars.Series``, or a
+``pandas.DataFrame`` or ``polars.DataFrame``. If a ``pandas.DataFrame`` or
+``polars.DataFrame`` is provided, the output is returned as the same type
+with named output columns.
 
 For example, inputs could be provided for the typical "OHLCV" data:
 

requirements.txt

@@ -1,2 +1,2 @@
-numpy>=1.19.4
-Cython>=0.29.21
+numpy
+Cython

requirements_dev.txt

@@ -1,4 +1,4 @@
 -r requirements.txt
-beautifulsoup4>=4.9.3
-mistune>=0.8.4
-Pygments>=2.7.4
+beautifulsoup4
+mistune
+Pygments

requirements_test.txt

@@ -1,3 +1,4 @@
 -r requirements.txt
-pandas>=1.1.4
-nose>=1.3.7
+pandas
+nose
+polars

talib/__init__.py

@@ -49,6 +49,39 @@ def wrapper(*args, **kwargs):
 
         return wrapper
 
+# If polars is available, wrap talib functions so that they support
+# polars.Series input
+try:
+    from polars import Series as _pl_Series
+except ImportError:
+    # polars not available, nothing to wrap
+    _polars_wrapper = lambda x: x
+else:
+    def _polars_wrapper(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Use Series' float64 values if pandas, else use values as passed
+            args = [arg.to_numpy().astype(float) if isinstance(arg, _pl_Series) else arg
+                    for arg in args]
+            kwargs = {k: v.to_numpy().astype(float) if isinstance(v, _pl_Series) else v
+                      for k, v in kwargs.items()}
+
+            result = func(*args, **kwargs)
+
+            # check to see if we got a streaming result
+            first_result = result[0] if isinstance(result, tuple) else result
+            is_streaming_fn_result = not hasattr(first_result, '__len__')
+            if is_streaming_fn_result:
+                return result
+
+            # Series was passed in, Series gets out
+            if isinstance(result, tuple):
+                # Handle multi-array results such as BBANDS
+                return tuple(_pl_Series(arr) for arr in result)
+            return _pl_Series(result)
+
+        return wrapper
+
 from ._ta_lib import (
     _ta_initialize, _ta_shutdown, MA_Type, __ta_version__,
     _ta_set_unstable_period as set_unstable_period,
@@ -75,6 +108,20 @@ def wrapper(*args, **kwargs):
     setattr(stream, func_name, wrapped_func)
     globals()[stream_func_name] = wrapped_func
 
+# wrap them with polars
+func = __import__("_ta_lib", globals(), locals(), __TA_FUNCTION_NAMES__, level=1)
+for func_name in __TA_FUNCTION_NAMES__:
+    wrapped_func = _polars_wrapper(getattr(func, func_name))
+    setattr(func, func_name, wrapped_func)
+    globals()[func_name] = wrapped_func
+
+stream_func_names = ['stream_%s' % fname for fname in __TA_FUNCTION_NAMES__]
+stream = __import__("stream", globals(), locals(), stream_func_names, level=1)
+for func_name, stream_func_name in zip(__TA_FUNCTION_NAMES__, stream_func_names):
+    wrapped_func = _polars_wrapper(getattr(stream, func_name))
+    setattr(stream, func_name, wrapped_func)
+    globals()[stream_func_name] = wrapped_func
+
 __version__ = '0.4.22'
 
 # In order to use this python library, talib (i.e. this __file__) will be

talib/_abstract.pxi

@@ -29,19 +29,35 @@ __INPUT_PRICE_SERIES_DEFAULTS = {'price':   'close',
                                  'periods': 'periods', # only used by MAVP; not a price series!
                                  }
 
+__INPUT_ARRAYS_TYPES = [dict]
+__ARRAY_TYPES = [np.ndarray]
+
 # allow use of pandas.DataFrame for input arrays
 try:
     import pandas
-    __INPUT_ARRAYS_TYPES = (dict, pandas.DataFrame)
-    __ARRAY_TYPES = (np.ndarray, pandas.Series)
+    __INPUT_ARRAYS_TYPES.append(pandas.DataFrame)
+    __ARRAY_TYPES.append(pandas.Series)
     __PANDAS_DATAFRAME = pandas.DataFrame
     __PANDAS_SERIES = pandas.Series
 except ImportError:
-    __INPUT_ARRAYS_TYPES = (dict,)
-    __ARRAY_TYPES = (np.ndarray,)
     __PANDAS_DATAFRAME = None
     __PANDAS_SERIES = None
 
+# allow use of polars.DataFrame for input arrays
+try:
+    import polars
+    __INPUT_ARRAYS_TYPES.append(polars.DataFrame)
+    __ARRAY_TYPES.append(polars.Series)
+    __POLARS_DATAFRAME = polars.DataFrame
+    __POLARS_SERIES = polars.Series
+except ImportError:
+    __POLARS_DATAFRAME = None
+    __POLARS_SERIES = None
+
+__INPUT_ARRAYS_TYPES = tuple(__INPUT_ARRAYS_TYPES)
+__ARRAY_TYPES = tuple(__ARRAY_TYPES)
+
+
 if sys.version >= '3':
 
     def str2bytes(s):
@@ -64,10 +80,10 @@ class Function(object):
     intended to simplify using individual TALIB functions by providing a
     unified interface for setting/controlling input data, setting function
     parameters and retrieving results. Input data consists of a ``dict`` of
-    ``numpy`` arrays (or a ``pandas.DataFrame``), one array for each of open,
-    high, low, close and volume. This can be set with the set_input_arrays()
-    method. Which keyed array(s) are used as inputs when calling the function
-    is controlled using the input_names property.
+    ``numpy`` arrays (or a ``pandas.DataFrame`` or ``polars.DataFrame``), one
+    array for each of open, high, low, close and volume. This can be set with
+    the set_input_arrays() method. Which keyed array(s) are used as inputs when
+    calling the function is controlled using the input_names property.
 
     This class gets initialized with a TALIB function name and optionally an
     input_arrays object. It provides the following primary functions for
@@ -334,6 +350,13 @@ class Function(object):
                 return __PANDAS_DATAFRAME(numpy.column_stack(ret),
                                           index=index,
                                           columns=self.output_names)
+        elif __POLARS_DATAFRAME is not None and \
+                isinstance(self.__input_arrays, __POLARS_DATAFRAME):
+            if len(ret) == 1:
+                return __POLARS_SERIES(ret[0])
+            else:
+                return __POLARS_DATAFRAME(numpy.column_stack(ret),
+                                          columns=self.output_names)
         else:
             return ret[0] if len(ret) == 1 else ret
 
@@ -382,6 +405,9 @@ class Function(object):
         if __PANDAS_DATAFRAME is not None \
                 and isinstance(self.__input_arrays, __PANDAS_DATAFRAME):
             no_existing_input_arrays = self.__input_arrays.empty
+        elif __POLARS_DATAFRAME is not None \
+                and isinstance(self.__input_arrays, __POLARS_DATAFRAME):
+            no_existing_input_arrays = self.__input_arrays.empty
         else:
             no_existing_input_arrays = not bool(self.__input_arrays)
 
@@ -432,6 +458,9 @@ class Function(object):
             if __PANDAS_SERIES is not None and \
                     isinstance(series, __PANDAS_SERIES):
                 series = series.values.astype(float)
+            elif __POLARS_SERIES is not None and \
+                    isinstance(series, __POLARS_SERIES):
+                series = series.to_numpy().astype(float)
             args.append(series)
         for opt_input in self.__opt_inputs:
             value = self.__get_opt_input_value(opt_input)

talib/test_polars.py

@@ -0,0 +1,41 @@
+import numpy as np
+import polars as pl
+from nose.tools import assert_equals, assert_is_instance, assert_true
+
+import talib
+from talib.test_data import series, assert_np_arrays_equal
+
+def test_MOM():
+    values = pl.Series([90.0,88.0,89.0])
+    result = talib.MOM(values, timeperiod=1)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan, -2, 1])
+    result = talib.MOM(values, timeperiod=2)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan, np.nan, -1])
+    result = talib.MOM(values, timeperiod=3)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan, np.nan, np.nan])
+    result = talib.MOM(values, timeperiod=4)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan, np.nan, np.nan])
+
+def test_MAVP():
+    a = pl.Series([1,5,3,4,7,3,8,1,4,6], dtype=pl.Float64)
+    b = pl.Series([2,4,2,4,2,4,2,4,2,4], dtype=pl.Float64)
+    result = talib.MAVP(a, b, minperiod=2, maxperiod=4)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan,np.nan,np.nan,3.25,5.5,4.25,5.5,4.75,2.5,4.75])
+    sma2 = talib.SMA(a, 2)
+    assert_is_instance(sma2, pl.Series)
+    assert_np_arrays_equal(result.to_numpy()[4::2], sma2.to_numpy()[4::2])
+    sma4 = talib.SMA(a, 4)
+    assert_is_instance(sma4, pl.Series)
+    assert_np_arrays_equal(result.to_numpy()[3::2], sma4.to_numpy()[3::2])
+    result = talib.MAVP(a, b, minperiod=2, maxperiod=3)
+    assert_is_instance(result, pl.Series)
+    assert_np_arrays_equal(result.to_numpy(), [np.nan,np.nan,4,4,5.5,4.666666666666667,5.5,4,2.5,3.6666666666666665])
+    sma3 = talib.SMA(a, 3)
+    assert_is_instance(sma3, pl.Series)
+    assert_np_arrays_equal(result.to_numpy()[2::2], sma2.to_numpy()[2::2])
+    assert_np_arrays_equal(result.to_numpy()[3::2], sma3.to_numpy()[3::2])