From 008ab547c71ab8ab3301c7eece31977b1b76a9b0 Mon Sep 17 00:00:00 2001
From: David Brazda <davidbrazda61@gmail.com>
Date: Wed, 30 Oct 2024 14:25:34 +0100
Subject: [PATCH] fix

---
 README.md         | 45 ++++++++++++++-------------------------------
 ttools/loaders.py | 10 ++++++++--
 2 files changed, 22 insertions(+), 33 deletions(-)

diff --git a/README.md b/README.md
index e049879..92b6b2b 100644
--- a/README.md
+++ b/README.md
@@ -21,40 +21,23 @@ Returns vectorized aggregation of given type. If trades for given period are not
 Example:
 
 ```python
+from ttools import load_data
 #This is how to call LOAD function
-symbol = ["BAC"]
-#datetime in zoneNY 
-day_start = datetime(2024, 10, 14, 9, 45, 0)
-day_stop = datetime(2024, 10, 16, 15, 1, 0)
-day_start = zoneNY.localize(day_start)
-day_stop = zoneNY.localize(day_stop)
-
-#requested AGG
-resolution = 12 #12s
-agg_type = AggType.OHLCV #other types AggType.OHLCV_VOL, AggType.OHLCV_DOL, AggType.OHLCV_RENKO
-exclude_conditions = ['C','O','4','B','7','V','P','W','U','Z','F','9','M','6'] #None to defaults
-minsize = 100
-main_session_only = True
-force_remote = False
-
-data = load_data(symbol = symbol,
-                     agg_type = agg_type,
-                     resolution = resolution,
-                     start_date = day_start,
-                     end_date = day_stop,
-                     #exclude_conditions = None,
-                     minsize = minsize,
-                     main_session_only = main_session_only,
-                     force_remote = force_remote,
-                     return_vbt = True, #returns vbt object
-                     verbose = True
+vbt_data = load_data(symbol = ["BAC"],
+                     agg_type = AggType.OHLCV, #aggregation types: AggType.OHLCV_VOL, AggType.OHLCV_DOL, AggType.OHLCV_RENKO,
+                     resolution = 12, #12s (for other types might be bricksize etc.)
+                     start_date = datetime(2024, 10, 14, 9, 45, 0),
+                     end_date = datetime(2024, 10, 16, 15, 1, 0),
+                     #exclude_conditions = ['C','O','4','B','7','V','P','W','U','Z','F','9','M','6'],
+                     minsize = 100, #minimum trade size included in aggregation
+                     main_session_only = True, #False for ext hours
+                     force_remote = False, #always refetches trades remotely
+                     return_vbt = True, #returns vbt object with symbols as columns, otherwise dict keyed by symbols with pd.DataFrame
+                     verbose = True # False = silent mode
                      )
-bac_df = ohlcv_df["BAC"]
 
-basic_data = vbt.Data.from_data(vbt.symbol_dict(ohlcv_df), tz_convert=zoneNY)
-vbt.settings['plotting']['auto_rangebreaks'] = True
-basic_data.ohlcv.plot()
-data.ohlcv.data[symbol[0]].lw.plot()
+vbt_data.ohlcv.data[symbol[0]].lw.plot()
+vbt_data.data[symbol[0]]
 ```
 ## prepare trade cache
 
diff --git a/ttools/loaders.py b/ttools/loaders.py
index f61a0e3..df4da15 100644
--- a/ttools/loaders.py
+++ b/ttools/loaders.py
@@ -359,8 +359,8 @@ def load_data(symbol: Union[str, List[str]],
         symbol (Union[str, list]): Symbol
         agg_type (AggType): Type of aggregation
         resolution (Union[str, int]) Resolution of aggregation nased on agg_type 
-        start_date (datetime): 
-        end_date (datetime): 
+        start_date (datetime): Start period, timezone aware
+        end_date (datetime): Start period, timezone aware
         exclude_conditions (list, optional): Trade conditions to exclude. Defaults to None.
         minsize (_type_, optional): Minimum trade size to include. Defaults to None.
         main_session_only (bool, optional): Main or ext. hours.. Defaults to True.
@@ -379,6 +379,12 @@ def load_data(symbol: Union[str, List[str]],
     if verbose is not None:
         set_verbose(verbose)  # Change global verbose if specified    
 
+    if start_date.tzinfo is None:
+        start_date = zoneNY.localize(start_date)
+
+    if end_date.tzinfo is None:
+        end_date = zoneNY.localize(end_date)
+
     if exclude_conditions is None:
         exclude_conditions = EXCLUDE_CONDITIONS