PERF: leverage tzlocal package to provide 2000x speedup for dateutil.tz.tzlocal operations

qwhelan · qwhelan · commit 88206af16e4d · 2019-01-11T21:57:32.000-08:00
diff --git a/asv_bench/asv.conf.json b/asv_bench/asv.conf.json
@@ -51,6 +51,7 @@
         "xlrd": [],
         "xlwt": [],
         "pytest": [],
+        "tzlocal": [],
         // If using Windows with python 2.7 and want to build using the
         // mingw toolchain (rather than MSVC), uncomment the following line.
         // "libpython": [],
diff --git a/ci/deps/azure-27-compat.yaml b/ci/deps/azure-27-compat.yaml
@@ -10,10 +10,11 @@ dependencies:
   - numpy=1.12.0
   - openpyxl=2.5.5
   - pytables=3.4.2
-  - python-dateutil=2.5.0
+  - python-dateutil=2.6.0
   - python=2.7*
   - pytz=2013b
   - scipy=0.18.1
+  - tzlocal=1.5.1
   - xlrd=1.0.0
   - xlsxwriter=0.5.2
   - xlwt=0.7.5
diff --git a/ci/deps/azure-27-locale.yaml b/ci/deps/azure-27-locale.yaml
@@ -16,6 +16,7 @@ dependencies:
   - pytz=2013b
   - scipy
   - sqlalchemy=0.8.1
+  - tzlocal=1.5.1
   - xlrd=1.0.0
   - xlsxwriter=0.5.2
   - xlwt=0.7.5
diff --git a/ci/deps/azure-macos-35.yaml b/ci/deps/azure-macos-35.yaml
@@ -17,6 +17,7 @@ dependencies:
   - pytables
   - python=3.5*
   - pytz
+  - tzlocal
   - xarray
   - xlrd
   - xlsxwriter
@@ -25,5 +26,5 @@ dependencies:
   - pytest
   - pytest-xdist
   - pip:
-    - python-dateutil==2.5.3
+    - python-dateutil==2.6.0
     - hypothesis>=3.58.0
diff --git a/ci/deps/azure-windows-27.yaml b/ci/deps/azure-windows-27.yaml
@@ -20,6 +20,7 @@ dependencies:
   - s3fs
   - scipy
   - sqlalchemy
+  - tzlocal
   - xlrd
   - xlsxwriter
   - xlwt
diff --git a/ci/deps/azure-windows-36.yaml b/ci/deps/azure-windows-36.yaml
@@ -18,6 +18,7 @@ dependencies:
   - python=3.6.6
   - pytz
   - scipy
+  - tzlocal
   - xlrd
   - xlsxwriter
   - xlwt
diff --git a/ci/deps/travis-27.yaml b/ci/deps/travis-27.yaml
@@ -28,12 +28,13 @@ dependencies:
   - pytables
   - blosc=1.14.3
   - python-blosc
-  - python-dateutil=2.5.0
+  - python-dateutil=2.6.0
   - python=2.7*
   - pytz=2013b
   - s3fs
   - scipy
   - sqlalchemy=0.9.6
+  - tzlocal
   - xarray=0.9.6
   - xlrd=1.0.0
   - xlsxwriter=0.5.2
diff --git a/ci/deps/travis-36-doc.yaml b/ci/deps/travis-36-doc.yaml
@@ -36,6 +36,7 @@ dependencies:
   - sphinx
   - sqlalchemy
   - statsmodels
+  - tzlocal
   - xarray
   - xlrd
   - xlsxwriter
diff --git a/doc/source/install.rst b/doc/source/install.rst
@@ -241,6 +241,8 @@ Recommended Dependencies
   evaluations. ``bottleneck`` uses specialized cython routines to achieve large speedups. If installed,
   must be Version 1.2.0 or higher.
 
+* `tzlocal <https://github.com/regebro/tzlocal>`__: significantly beneficial to performance if using the `dateutil.tz.tzlocal()` timezone.
+
 .. note::
 
    You are highly encouraged to install these libraries, as they provide speed improvements, especially
@@ -328,7 +330,6 @@ Optional Dependencies
        to get the necessary dependencies for installation of `lxml`_. This
        will prevent further headaches down the line.
 
-
 .. _html5lib: https://github.com/html5lib/html5lib-python
 .. _BeautifulSoup4: http://www.crummy.com/software/BeautifulSoup
 .. _lxml: http://lxml.de
diff --git a/environment.yml b/environment.yml
@@ -6,7 +6,7 @@ dependencies:
   # required
   - numpy>=1.15
   - python=3
-  - python-dateutil>=2.5.0
+  - python-dateutil>=2.6.0
   - pytz
 
   # development
@@ -48,6 +48,7 @@ dependencies:
   - seaborn
   - sqlalchemy
   - statsmodels
+  - tzlocal
   - xarray
   - xlrd
   - xlsxwriter
diff --git a/pandas/_libs/tslib.pyx b/pandas/_libs/tslib.pyx
@@ -28,7 +28,8 @@ from pandas._libs.tslibs.np_datetime import OutOfBoundsDatetime
 from pandas._libs.tslibs.parsing import parse_datetime_string
 
 from pandas._libs.tslibs.timedeltas cimport cast_from_unit
-from pandas._libs.tslibs.timezones cimport is_utc, is_tzlocal, get_dst_info
+from pandas._libs.tslibs.timezones cimport (is_utc, is_tzlocal, get_dst_info,
+    get_tzlocal_tz)
 from pandas._libs.tslibs.timezones import UTC
 from pandas._libs.tslibs.conversion cimport (
     tz_convert_single, _TSObject, convert_datetime_to_tsobject,
@@ -124,6 +125,9 @@ def ints_to_pydatetime(int64_t[:] arr, object tz=None, object freq=None,
         raise ValueError("box must be one of 'datetime', 'date', 'time' or"
                          " 'timestamp'")
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if is_utc(tz) or tz is None:
         for i in range(n):
             value = arr[i]
diff --git a/pandas/_libs/tslibs/conversion.pyx b/pandas/_libs/tslibs/conversion.pyx
@@ -33,7 +33,7 @@ from pandas._libs.tslibs.timedeltas cimport (cast_from_unit,
                                              delta_to_nanoseconds)
 from pandas._libs.tslibs.timezones cimport (
     is_utc, is_tzlocal, is_fixed_offset, get_utcoffset, get_dst_info,
-    get_timezone, maybe_get_tz, tz_compare)
+    get_timezone, maybe_get_tz, tz_compare, get_tzlocal_tz)
 from pandas._libs.tslibs.timezones import UTC
 from pandas._libs.tslibs.parsing import parse_datetime_string
 
@@ -541,9 +541,19 @@ cdef inline void localize_tso(_TSObject obj, tzinfo tz):
         int64_t local_val
         Py_ssize_t pos
         str typ
+        bint tz_changed
 
     assert obj.tzinfo is None
 
+    orig_tz = tz
+    tz_changed = False
+
+    if is_tzlocal(tz):
+        new_tz = get_tzlocal_tz(tz)
+        if new_tz != tz:
+            tz_changed = True
+            tz = new_tz
+
     if is_utc(tz):
         pass
     elif obj.value == NPY_NAT:
@@ -574,7 +584,12 @@ cdef inline void localize_tso(_TSObject obj, tzinfo tz):
             # so this branch will never be reached.
             pass
 
-    obj.tzinfo = tz
+    if tz_changed:
+        # We want to return tzlocal() if provided it, even if we map it
+        # to a real tz for performance reasons
+        obj.tzinfo = orig_tz
+    else:
+        obj.tzinfo = tz
 
 
 cdef inline datetime _localize_pydatetime(datetime dt, tzinfo tz):
@@ -652,6 +667,9 @@ cdef inline int64_t[:] _tz_convert_dst(int64_t[:] values, tzinfo tz,
         bint tz_is_local
 
     tz_is_local = is_tzlocal(tz)
+    if tz_is_local:
+        tz = get_tzlocal_tz(tz)
+        tz_is_local = is_tzlocal(tz)
 
     if not tz_is_local:
         # get_dst_info cannot extract offsets from tzlocal because its
@@ -763,6 +781,11 @@ cpdef int64_t tz_convert_single(int64_t val, object tz1, object tz2):
     if val == NPY_NAT:
         return val
 
+    if is_tzlocal(tz1):
+        tz1 = get_tzlocal_tz(tz1)
+    if is_tzlocal(tz2):
+        tz2 = get_tzlocal_tz(tz2)
+
     # Convert to UTC
     if is_tzlocal(tz1):
         utc_date = _tz_convert_tzlocal_utc(val, tz1, to_utc=True)
@@ -807,10 +830,16 @@ cdef inline int64_t[:] _tz_convert_one_way(int64_t[:] vals, object tz,
         int64_t[:] converted, result
         Py_ssize_t i, n = len(vals)
         int64_t val
+        bint tz_is_local
+
+    tz_is_local = is_tzlocal(tz)
+    if tz_is_local:
+        tz = get_tzlocal_tz(tz)
+        tz_is_local = is_tzlocal(tz)
 
     if not is_utc(get_timezone(tz)):
         converted = np.empty(n, dtype=np.int64)
-        if is_tzlocal(tz):
+        if tz_is_local:
             for i in range(n):
                 val = vals[i]
                 if val == NPY_NAT:
@@ -915,6 +944,9 @@ def tz_localize_to_utc(ndarray[int64_t] vals, object tz, object ambiguous=None,
 
     result = np.empty(n, dtype=np.int64)
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if is_tzlocal(tz):
         for i in range(n):
             v = vals[i]
@@ -1222,6 +1254,9 @@ cdef int64_t[:] _normalize_local(int64_t[:] stamps, tzinfo tz):
         npy_datetimestruct dts
         int64_t delta, local_val
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if is_tzlocal(tz):
         for i in range(n):
             if stamps[i] == NPY_NAT:
@@ -1301,6 +1336,9 @@ def is_date_array_normalized(int64_t[:] stamps, object tz=None):
         int64_t local_val, delta
         str typ
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if tz is None or is_utc(tz):
         for i in range(n):
             dt64_to_dtstruct(stamps[i], &dts)
diff --git a/pandas/_libs/tslibs/period.pyx b/pandas/_libs/tslibs/period.pyx
@@ -32,7 +32,8 @@ cimport pandas._libs.tslibs.util as util
 from pandas._libs.tslibs.util cimport is_period_object, is_string_object
 
 from pandas._libs.tslibs.timestamps import Timestamp
-from pandas._libs.tslibs.timezones cimport is_utc, is_tzlocal, get_dst_info
+from pandas._libs.tslibs.timezones cimport (is_utc, is_tzlocal, get_dst_info,
+    get_tzlocal_tz)
 from pandas._libs.tslibs.timedeltas import Timedelta
 from pandas._libs.tslibs.timedeltas cimport delta_to_nanoseconds
 
@@ -1505,6 +1506,9 @@ cdef int64_t[:] localize_dt64arr_to_period(int64_t[:] stamps,
         npy_datetimestruct dts
         int64_t local_val
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if is_utc(tz) or tz is None:
         with nogil:
             for i in range(n):
diff --git a/pandas/_libs/tslibs/resolution.pyx b/pandas/_libs/tslibs/resolution.pyx
@@ -9,7 +9,7 @@ from pandas._libs.tslibs.np_datetime cimport (
     npy_datetimestruct, dt64_to_dtstruct)
 from pandas._libs.tslibs.frequencies cimport get_freq_code
 from pandas._libs.tslibs.timezones cimport (
-    is_utc, is_tzlocal, maybe_get_tz, get_dst_info)
+    is_utc, is_tzlocal, maybe_get_tz, get_dst_info, get_tzlocal_tz)
 from pandas._libs.tslibs.conversion cimport tz_convert_utc_to_tzlocal
 from pandas._libs.tslibs.ccalendar cimport get_days_in_month
 
@@ -49,6 +49,9 @@ cdef _reso_local(int64_t[:] stamps, object tz):
         npy_datetimestruct dts
         int64_t local_val, delta
 
+    if is_tzlocal(tz):
+        tz = get_tzlocal_tz(tz)
+
     if is_utc(tz) or tz is None:
         for i in range(n):
             if stamps[i] == NPY_NAT:
diff --git a/pandas/_libs/tslibs/timezones.pxd b/pandas/_libs/tslibs/timezones.pxd
@@ -3,6 +3,10 @@
 cpdef bint is_utc(object tz)
 cdef bint is_tzlocal(object tz)
 
+cpdef bint have_tzlocal_package()
+cpdef object get_tzlocal_tz(object tz)
+cpdef _set_tzlocal_tz(object tz)
+
 cdef bint treat_tz_as_pytz(object tz)
 cdef bint treat_tz_as_dateutil(object tz)
 
diff --git a/pandas/_libs/tslibs/timezones.pyx b/pandas/_libs/tslibs/timezones.pyx
@@ -12,6 +12,8 @@ from pytz.tzinfo import BaseTzInfo as _pytz_BaseTzInfo
 import pytz
 UTC = pytz.utc
 
+tzlocal_package = None
+tzlocal_tz = None
 
 import numpy as np
 cimport numpy as cnp
@@ -34,6 +36,38 @@ cdef inline bint is_tzlocal(object tz):
     return isinstance(tz, _dateutil_tzlocal)
 
 
+cpdef bint have_tzlocal_package():
+    global tzlocal_package
+    if tzlocal_package is None:
+        try:
+            import tzlocal
+            tzlocal_package = tzlocal
+        except ImportError:
+            pass
+
+    return tzlocal_package is not None
+
+
+cpdef object get_tzlocal_tz(object tz):
+    if have_tzlocal_package():
+        global tzlocal_tz
+        if tzlocal_tz is None:
+            local_tz = tzlocal_package.get_localzone()
+            tzlocal_tz = local_tz.zone
+
+        try:
+            return pytz.timezone(tzlocal_tz)
+        except pytz.exceptions.UnknownTimeZoneError:
+            tzlocal_tz = None
+
+    return tz
+
+
+cpdef _set_tzlocal_tz(object tz_str):
+    global tzlocal_tz
+    tzlocal_tz = str(tz_str)
+
+
 cdef inline bint treat_tz_as_pytz(object tz):
     return (hasattr(tz, '_utc_transition_times') and
             hasattr(tz, '_transition_info'))
diff --git a/pandas/tests/indexes/datetimes/test_timezones.py b/pandas/tests/indexes/datetimes/test_timezones.py
@@ -820,7 +820,8 @@ def test_timetz_accessor(self, tz_naive_fixture):
 
         index = DatetimeIndex(['2018-06-04 10:20:30', pd.NaT], tz=tz)
         result = index.timetz
-
+        print('result', result)
+        print('expected', expected)
         tm.assert_numpy_array_equal(result, expected)
 
     def test_dti_drop_dont_lose_tz(self):
diff --git a/pandas/tests/series/test_datetime_values.py b/pandas/tests/series/test_datetime_values.py
@@ -536,6 +536,8 @@ def test_dt_timetz_accessor(self, tz_naive_fixture):
         expected = Series([time(23, 56, tzinfo=tz), time(21, 24, tzinfo=tz),
                            time(22, 14, tzinfo=tz)])
         result = s.dt.timetz
+        print('result', result)
+        print('expected', expected)
         tm.assert_series_equal(result, expected)
 
     def test_setitem_with_string_index(self):
diff --git a/pandas/tests/tslibs/test_timezones.py b/pandas/tests/tslibs/test_timezones.py
@@ -49,6 +49,30 @@ def test_tzlocal_offset():
     assert ts.value + offset == Timestamp("2011-01-01").value
 
 
+def test_tzlocal_package():
+    if not timezones.have_tzlocal_package():
+        pytest.skip("tzlocal package is not installed, can't test behavior")
+
+    orig_tzlocal_tz = timezones.get_tzlocal_tz(None)
+    # Check that we fall back to non-tzlocal code path if user has an invalid
+    # system tz set
+    timezones._set_tzlocal_tz('foo/bar')
+    ts = Timestamp("2011-01-01", tz=dateutil.tz.tzlocal())
+    assert ts.tz == dateutil.tz.tzlocal()
+
+    utc = pytz.utc
+    for tz in pytz.all_timezones:
+        # Check that all pytz timezones work
+        ts_tz = Timestamp("2011-01-01", tz=tz)
+        timezones._set_tzlocal_tz(tz)
+        ts_local = Timestamp("2011-01-01", tz=dateutil.tz.tzlocal())
+
+        assert ts_tz.astimezone(utc) == ts_local.astimezone(utc)
+        assert ts_local.tz == dateutil.tz.tzlocal()
+
+    timezones._set_tzlocal_tz(orig_tzlocal_tz.zone)
+
+
 @pytest.fixture(params=[
     (pytz.timezone("US/Eastern"), lambda tz, x: tz.localize(x)),
     (dateutil.tz.gettz("US/Eastern"), lambda tz, x: x.replace(tzinfo=tz))
diff --git a/requirements-dev.txt b/requirements-dev.txt
diff --git a/setup.py b/setup.py
