From 201b4959b865f8ab470f88be5edeb289cfad3021 Mon Sep 17 00:00:00 2001
From: Quyen Nguyen <122257324+qnguyen345@users.noreply.github.com>
Date: Fri, 9 Feb 2024 06:58:06 -0700
Subject: [PATCH] Allow tz-aware timstamps in `gaps.trim` (#206)

* updated gaps file for later versions of pandas

* fixed tz issues

* included new note for gaps function

* included notes and unit testing for gaps.trim

* Fixed pep8 format

* deleted .date()

---------

Co-authored-by: Nguyen <qnguyen@nrel.gov>
---
 docs/whatsnew/0.2.0.rst                |  3 +++
 pvanalytics/quality/gaps.py            |  3 ++-
 pvanalytics/tests/quality/test_gaps.py | 21 +++++++++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/docs/whatsnew/0.2.0.rst b/docs/whatsnew/0.2.0.rst
index 9d84a927..a00116ad 100644
--- a/docs/whatsnew/0.2.0.rst
+++ b/docs/whatsnew/0.2.0.rst
@@ -40,6 +40,8 @@ Bug Fixes
   of raising ``AttributeError``. (:pull:`181`)
 * Compatibility with pandas 2.0.0 (:pull:`185`) and future versions of pandas (:pull:`203`)
 * Compatibility with scipy 1.11 (:pull:`196`)
+* Updated function :py:func:`~pvanalytics.quality.gaps.trim` to handle pandas 2.0.0 update for tz-aware timeseries
+ (:pull:`206`)
 
 Requirements
 ~~~~~~~~~~~~
@@ -64,3 +66,4 @@ Contributors
 * Kevin Anderson (:ghuser:`kanderso-nrel`)
 * Cliff Hansen (:ghuser:`cwhanse`)
 * Abhishek Parikh (:ghuser:`abhisheksparikh`)
+* Quyen Nguyen (:ghuser:`qnguyen345`)
diff --git a/pvanalytics/quality/gaps.py b/pvanalytics/quality/gaps.py
index e54c381a..642a2db5 100644
--- a/pvanalytics/quality/gaps.py
+++ b/pvanalytics/quality/gaps.py
@@ -409,8 +409,9 @@ def trim(series, days=10):
     """
     start, end = start_stop_dates(series, days=days)
     mask = pd.Series(False, index=series.index)
+
     if start:
-        mask.loc[start.date():end.date()] = True
+        mask.loc[start:end] = True
     return mask
 
 
diff --git a/pvanalytics/tests/quality/test_gaps.py b/pvanalytics/tests/quality/test_gaps.py
index a5a06799..51cce2e9 100644
--- a/pvanalytics/tests/quality/test_gaps.py
+++ b/pvanalytics/tests/quality/test_gaps.py
@@ -449,6 +449,27 @@ def test_trim_daily_index():
     )
 
 
+def test_trim_daily_index_tz_aware():
+    """trim works when data has a daily index and data is tz-aware."""
+    data = pd.Series(True, index=pd.date_range(
+        start='1/1/2020', end='2/29/2020', freq='D', tz="Etc/GMT+8"))
+    assert gaps.trim(data).all()
+    data.iloc[0:8] = False
+    data.iloc[9] = False
+    expected = data.copy()
+    expected.iloc[0:10] = False
+    assert_series_equal(
+        expected,
+        gaps.trim(data)
+    )
+    data.iloc[-5:] = False
+    expected.iloc[-5:] = False
+    assert_series_equal(
+        expected,
+        gaps.trim(data)
+    )
+
+
 def test_completeness_score_all_nans():
     """A data set with all nans has completeness 0 for each day."""
     completeness = gaps.completeness_score(