From e4b7e0c11001188cfadc7cd36bde6e59b8eeaba8 Mon Sep 17 00:00:00 2001
From: jbrockmendel <jbrockmendel@gmail.com>
Date: Sun, 4 Aug 2019 09:25:50 -0700
Subject: [PATCH 1/2] BUG: Fix NaT +/- DTA/TDA

---
 pandas/_libs/tslibs/nattype.pyx |  9 ++++++---
 pandas/tests/scalar/test_nat.py | 19 +++++++++++++------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/pandas/_libs/tslibs/nattype.pyx b/pandas/_libs/tslibs/nattype.pyx
index 7f35a11e57b71..6fab1b5c02be1 100644
--- a/pandas/_libs/tslibs/nattype.pyx
+++ b/pandas/_libs/tslibs/nattype.pyx
@@ -123,7 +123,9 @@ cdef class _NaT(datetime):
             return c_NaT
         elif getattr(other, '_typ', None) in ['dateoffset', 'series',
                                               'period', 'datetimeindex',
-                                              'timedeltaindex']:
+                                              'datetimearray',
+                                              'timedeltaindex',
+                                              'timedeltaarray']:
             # Duplicate logic in _Timestamp.__add__ to avoid needing
             # to subclass; allows us to @final(_Timestamp.__add__)
             return NotImplemented
@@ -151,9 +153,10 @@ cdef class _NaT(datetime):
             return self + neg_other
 
         elif getattr(other, '_typ', None) in ['period', 'series',
-                                              'periodindex', 'dateoffset']:
+                                              'periodindex', 'dateoffset',
+                                              'datetimearray',
+                                              'timedeltaarray']:
             return NotImplemented
-
         return NaT
 
     def __pos__(self):
diff --git a/pandas/tests/scalar/test_nat.py b/pandas/tests/scalar/test_nat.py
index f935a7fa880c7..576569a8a0953 100644
--- a/pandas/tests/scalar/test_nat.py
+++ b/pandas/tests/scalar/test_nat.py
@@ -18,7 +18,7 @@
     Timestamp,
     isna,
 )
-from pandas.core.arrays import PeriodArray
+from pandas.core.arrays import DatetimeArray, PeriodArray, TimedeltaArray
 from pandas.util import testing as tm
 
 
@@ -397,7 +397,9 @@ def test_nat_rfloordiv_timedelta(val, expected):
     "value",
     [
         DatetimeIndex(["2011-01-01", "2011-01-02"], name="x"),
-        DatetimeIndex(["2011-01-01", "2011-01-02"], name="x"),
+        DatetimeIndex(["2011-01-01", "2011-01-02"], tz="US/Eastern", name="x"),
+        DatetimeArray._from_sequence(["2011-01-01", "2011-01-02"]),
+        DatetimeArray._from_sequence(["2011-01-01", "2011-01-02"], tz="US/Pacific"),
         TimedeltaIndex(["1 day", "2 day"], name="x"),
     ],
 )
@@ -406,19 +408,24 @@ def test_nat_arithmetic_index(op_name, value):
     exp_name = "x"
     exp_data = [NaT] * 2
 
-    if isinstance(value, DatetimeIndex) and "plus" in op_name:
-        expected = DatetimeIndex(exp_data, name=exp_name, tz=value.tz)
+    if value.dtype.kind == "M" and "plus" in op_name:
+        expected = DatetimeIndex(exp_data, tz=value.tz, name=exp_name)
     else:
         expected = TimedeltaIndex(exp_data, name=exp_name)
 
-    tm.assert_index_equal(_ops[op_name](NaT, value), expected)
+    if not isinstance(value, Index):
+        expected = expected.array
+
+    op = _ops[op_name]
+    result = op(NaT, value)
+    tm.assert_equal(result, expected)
 
 
 @pytest.mark.parametrize(
     "op_name",
     ["left_plus_right", "right_plus_left", "left_minus_right", "right_minus_left"],
 )
-@pytest.mark.parametrize("box", [TimedeltaIndex, Series])
+@pytest.mark.parametrize("box", [TimedeltaIndex, Series, TimedeltaArray._from_sequence])
 def test_nat_arithmetic_td64_vector(op_name, box):
     # see gh-19124
     vec = box(["1 day", "2 day"], dtype="timedelta64[ns]")

From 790df917d6a4b5a31a10e9add0fe854ea3e24900 Mon Sep 17 00:00:00 2001
From: jbrockmendel <jbrockmendel@gmail.com>
Date: Sun, 4 Aug 2019 14:25:16 -0700
Subject: [PATCH 2/2] requested edit to dtype check

---
 pandas/tests/scalar/test_nat.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pandas/tests/scalar/test_nat.py b/pandas/tests/scalar/test_nat.py
index 576569a8a0953..e7ad76cf95ba0 100644
--- a/pandas/tests/scalar/test_nat.py
+++ b/pandas/tests/scalar/test_nat.py
@@ -7,6 +7,8 @@
 from pandas._libs.tslibs import iNaT
 import pandas.compat as compat
 
+from pandas.core.dtypes.common import is_datetime64_any_dtype
+
 from pandas import (
     DatetimeIndex,
     Index,
@@ -408,7 +410,7 @@ def test_nat_arithmetic_index(op_name, value):
     exp_name = "x"
     exp_data = [NaT] * 2
 
-    if value.dtype.kind == "M" and "plus" in op_name:
+    if is_datetime64_any_dtype(value.dtype) and "plus" in op_name:
         expected = DatetimeIndex(exp_data, tz=value.tz, name=exp_name)
     else:
         expected = TimedeltaIndex(exp_data, name=exp_name)