MAINT test_strategy updated and process pools modified

pandas_ta/core.py

@@ -126,7 +126,7 @@ class BasePandasObject(PandasObject):
             # Preemptively drop the rows that are all NaNs
             # Might need to be moved to AnalysisIndicators.__call__() to be
             #   toggleable via kwargs.
-            # df.dropna(axis=0, inplace=True)
+            df.dropna(axis=0, inplace=True)
             # Preemptively rename columns to lowercase
             df.rename(columns=common_names, errors="ignore", inplace=True)
 
@@ -297,7 +297,7 @@ class AnalysisIndicators(BasePandasObject):
         else:
             self._mp = False
 
-    # Public Get DataFrame Methods
+    # Public Get DataFrame Properties
     @property
     def categories(self) -> str:
         """Returns the categories."""
@@ -364,11 +364,6 @@ class AnalysisIndicators(BasePandasObject):
         df = self._df
         if df is None: return
 
-        # def _get_case(column: str):
-        #     cases = [column.lower(), column.upper(), column.title()]
-        #     return [c for i, c in enumerate(cases) if column == cases[i]].pop()
-        # default = _get_case(default)
-
         # Explicitly passing a pd.Series to override default.
         if isinstance(series, pd.Series):
             return series
@@ -597,22 +592,14 @@ class AnalysisIndicators(BasePandasObject):
         if timed: stime = perf_counter()
         if mode["custom"]:
             # Custom multiprocessing pool. Must be ordered for Chained Strategies
-            results = pool.map(
+            results = pool.imap(
                 self._mp_worker,
                 [(ind["kind"], ind["params"] if "params" in ind and isinstance(ind["params"], tuple) else (), {**ind, **kwargs}) for ind in ta],
                 self.cores
             )
-
-            # # Without multiprocessing :
-            # for ind in ta:
-            #     params = ind["params"] if "params" in ind and isinstance(ind["params"], tuple) else tuple()
-            #     result = getattr(self, ind["kind"])(*params, **{**ind, **kwargs})
-            #     results.append(result)
-            #     self._add_prefix_suffix(result=result, **ind)
-            #     self._append(result=result, **kwargs)
         else:
-            # All and Categorical multiprocessing pool. Speed over Order
-            results = pool.map(
+            # All and Categorical multiprocessing pool. Speed over Order.
+            results = pool.imap_unordered(
                 self._mp_worker,
                 [(ind, tuple(), kwargs) for ind in ta],
                 self.cores

tests/test_strategy.py

@@ -1,20 +1,34 @@
+# Must run seperately from the rest of the tests
+# in order to successfully run
+from time import perf_counter
+
 from .config import sample_data
 from .context import pandas_ta
 
 from unittest import skip, TestCase
 from pandas import DataFrame
 
-# Must run seperately from the rest of the tests
-# in order to successfully run
+from pandas_ta.utils import final_time
+
+_verbose = False
+_timed = True
+speed_table = False
 
 class TestStrategyMethods(TestCase):
     @classmethod
     def setUpClass(cls):
         cls.data = sample_data
+        cls.speed_test = DataFrame()
 
     @classmethod
     def tearDownClass(cls):
         del cls.data
+        cls.speed_test = cls.speed_test.T
+        cls.speed_test.index.name = "Test"
+        cls.speed_test.columns = ["secs"]
+        cls.speed_test["Cumsecs"] = cls.speed_test["secs"].cumsum()
+        if speed_table: cls.speed_test.to_csv("tests/speed_test.csv")
+        print(cls.speed_test)
 
 
     def setUp(self): pass
@@ -22,8 +36,12 @@ class TestStrategyMethods(TestCase):
 
 
     def test_all(self):
+        if _verbose: print()
+        category = "All"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy(verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -31,9 +49,12 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_all_strategy(self):
+        if _verbose: print()
         init_cols = len(self.data.columns)
-        self.data.ta.strategy(pandas_ta.AllStrategy, verbose=False)
+        self.data.ta.strategy(pandas_ta.AllStrategy, verbose=_verbose)
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -42,8 +63,9 @@ class TestStrategyMethods(TestCase):
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
     def test_all_name_strategy(self):
+        if _verbose: print()
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("All", verbose=False)
+        self.data.ta.strategy("All", verbose=_verbose)
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -52,8 +74,12 @@ class TestStrategyMethods(TestCase):
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
     def test_candles_category(self):
+        if _verbose: print()
+        category = "Candles"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Candles", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -61,9 +87,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_common(self):
+        if _verbose: print()
+        category = "Common"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy(pandas_ta.CommonStrategy, verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(pandas_ta.CommonStrategy, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -72,6 +104,7 @@ class TestStrategyMethods(TestCase):
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
     def test_custom_a(self):
+        if _verbose: print()
         momo_bands_sma_ta = [
             {"kind":"sma", "length": 50}, # 1
             {"kind":"sma", "length": 200}, # 1
@@ -88,8 +121,12 @@ class TestStrategyMethods(TestCase):
             "MACD and RSI Momo with BBANDS and SMAs 50 & 200 and Cumulative Log Returns" # description
         )
 
+        category = "Custom A"
+
         init_cols = len(self.data.columns)
-        self.data.ta.strategy(custom, verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(custom, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertEqual(added_cols, 11)
 
@@ -97,7 +134,10 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_custom_args_tuple(self):
+        if _verbose: print()
         custom_args_ta = [
             {"kind":"fisher", "params": (13, 7)},
             {"kind":"macd", "params": (9, 19, 7)},
@@ -110,17 +150,27 @@ class TestStrategyMethods(TestCase):
             "Allow for easy filling in indicator arguments without naming them"
         )
 
+        category = "Custom B"
+
         init_cols = len(self.data.columns)
-        self.data.ta.strategy(custom, verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(custom, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
 
         result = self.data[self.data.columns[-added_cols:]]
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_momentum_category(self):
+        if _verbose: print()
+        category = "Momentum"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Momentum", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -128,9 +178,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_overlap_category(self):
+        if _verbose: print()
+        category = "Overlap"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Overlap", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -138,9 +194,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_performance_category(self):
+        if _verbose: print()
+        category = "Performance"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Performance", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -148,9 +210,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_statistics_category(self):
+        if _verbose: print()
+        category = "Statistics"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Statistics", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -158,9 +226,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_trend_category(self):
+        if _verbose: print()
+        category = "Trend"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Trend", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -168,9 +242,15 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_volatility_category(self):
+        if _verbose: print()
+        category = "Volatility"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Volatility", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
@@ -178,12 +258,20 @@ class TestStrategyMethods(TestCase):
         self.assertIsInstance(result, DataFrame)
         self.data.drop(columns=result.columns, axis=1, inplace=True)
 
+        self.speed_test[category] = [time_diff]
+
     def test_volume_category(self):
+        if _verbose: print()
+        category = "Volume"
         init_cols = len(self.data.columns)
-        self.data.ta.strategy("Volume", verbose=False)
+        if _timed: stime = perf_counter()
+        self.data.ta.strategy(category, verbose=_verbose)
+        if _timed: time_diff = perf_counter() - stime
         added_cols = len(self.data.columns) - init_cols
         self.assertGreaterEqual(added_cols, 1)
 
         result = self.data[self.data.columns[-added_cols:]]
         self.assertIsInstance(result, DataFrame)
-        self.data.drop(columns=result.columns, axis=1, inplace=True)
+        self.data.drop(columns=result.columns, axis=1, inplace=True)
+
+        self.speed_test[category] = [time_diff]