diff --git a/README.md b/README.md index 16a6051..52333c4 100644 --- a/README.md +++ b/README.md @@ -642,6 +642,7 @@ result = ta.cagr(df.close) * **Moving Average Choices**: dema, ema, fwma, hma, linreg, midpoint, pwma, rma, sinwma, sma, swma, t3, tema, trima, vidya, wma, zlma. * An _experimental_ and independent __Watchlist__ Class located in the [Examples](https://github.com/twopirllc/pandas-ta/tree/master/examples/watchlist.py) Directory that can be used in conjunction with the new __Strategy__ Class. * _Linear Regression_ (**linear_regression**) is a new utility method for Simple Linear Regression using _Numpy_ or _Scikit Learn_'s implementation. +* Added utility/convience function, ```to_utc```, to convert the DataFrame index to UTC. See: ```help(ta.to_utc)```
diff --git a/pandas_ta/utils/_time.py b/pandas_ta/utils/_time.py index f6ba803..1411ec8 100644 --- a/pandas_ta/utils/_time.py +++ b/pandas_ta/utils/_time.py @@ -1,13 +1,14 @@ # -*- coding: utf-8 -*- from datetime import datetime from time import localtime, perf_counter +from typing import Tuple -from pandas import DataFrame, date_range, Series, Timestamp +from pandas import DataFrame, Series, Timestamp from pandas_ta import EXCHANGE_TZ -def df_dates(df: DataFrame, dates: (str, list) = None) -> DataFrame: +def df_dates(df: DataFrame, dates: Tuple[str, list] = None) -> DataFrame: """Yields the DataFrame with the given dates""" if dates is None: return None if not isinstance(dates, list): @@ -41,7 +42,7 @@ def final_time(stime): return f"{time_diff * 1000:2.4f} ms ({time_diff:2.4f} s)" -def get_time(exchange: str = "NYSE", full:bool = True, to_string:bool = False) -> (None, str): +def get_time(exchange: str = "NYSE", full:bool = True, to_string:bool = False) -> Tuple[None, str]: """Returns Current Time, Day of the Year and Percentage, and the current time of the selected Exchange.""" tz = EXCHANGE_TZ["NYSE"] # Default is NYSE (Eastern Time Zone) @@ -88,6 +89,19 @@ def total_time(series: Series, tf: str = "years") -> float: return TimeFrame[tf] return TimeFrame["years"] + +def to_utc(df: DataFrame) -> DataFrame: + """Either localizes the DataFrame Index to UTC or it applies + tz_convert to set the Index to UTC. + """ + if not df.empty: + try: + df.index = df.index.tz_localize("UTC") + except TypeError: + df.index = df.index.tz_convert("UTC") + return df + + # Aliases mtd_df = df_month_to_date qtd_df = df_quarter_to_date diff --git a/setup.py b/setup.py index 867c740..b7fe038 100644 --- a/setup.py +++ b/setup.py @@ -17,7 +17,7 @@ setup( "pandas_ta.volatility", "pandas_ta.volume" ], - version=".".join(("0", "2", "30b")), + version=".".join(("0", "2", "31b")), description=long_description, long_description=long_description, author="Kevin Johnson", diff --git a/tests/test_utils.py b/tests/test_utils.py index f4a051b..ec70e8a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -7,6 +7,8 @@ from unittest.mock import patch import numpy as np import numpy.testing as npt from pandas import DataFrame, Series +from pandas.api.types import is_datetime64_ns_dtype, is_datetime64tz_dtype + data = { "zero": [0, 0], @@ -259,6 +261,11 @@ class TestUtilities(TestCase): self.assertEqual(self.utils.get_offset(-1.1), 0) self.assertEqual(self.utils.get_offset(1), 1) + def test_to_utc(self): + result = self.utils.to_utc(self.data.copy()) + self.assertTrue(is_datetime64_ns_dtype(result.index)) + self.assertTrue(is_datetime64tz_dtype(result.index)) + def test_total_time(self): result = self.utils.total_time(self.data) self.assertEqual(20.824093086926762, result)