From 5f49fa22cb52bb139a5d5c4ff05f1b12884fe23f Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Fri, 22 Jan 2016 15:00:04 -0500
Subject: [PATCH] MAINT: Upgrade numpy and fix warnings.

Mostly fixes ambiguous calls to numpy.full, and uses explicitly-united
NaT values.
---
 etc/requirements.txt                        |   4 +-
 setup.py                                    | 130 ++++++++++----------
 tests/pipeline/test_earnings.py             |   4 +-
 tests/pipeline/test_engine.py               |  47 ++++---
 tests/pipeline/test_factor.py               |  11 +-
 tests/pipeline/test_numerical_expression.py |  20 +--
 tests/pipeline/test_term.py                 |   1 +
 tests/test_examples.py                      |   7 +-
 zipline/data/us_equity_pricing.py           |   2 +-
 zipline/examples/pairtrade.py               |   4 +-
 zipline/pipeline/factors/events.py          |   6 +-
 zipline/pipeline/factors/technical.py       |   2 +-
 zipline/pipeline/loaders/utils.py           |   8 +-
 zipline/utils/factory.py                    |   3 +-
 zipline/utils/numpy_utils.py                |  21 +++-
 15 files changed, 149 insertions(+), 121 deletions(-)

diff --git a/etc/requirements.txt b/etc/requirements.txt
index 3f9f0456..ccf6359e 100644
--- a/etc/requirements.txt
+++ b/etc/requirements.txt
@@ -24,7 +24,7 @@ six==1.9.0
 # For fetching remote data
 requests==2.9.1
 
-Cython==0.22.1
+Cython==0.23.4
 
 # faster OrderedDict
 cyordereddict==0.2.2
@@ -41,7 +41,7 @@ networkx==1.9.1
 numexpr==2.4.3
 
 # On disk storage format for pipeline data.
-bcolz==0.10.0
+bcolz==0.12.1
 
 # Command line interface helper
 click==4.0.0
diff --git a/setup.py b/setup.py
index da09b60c..35788c60 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@ from os.path import (
     join,
 )
 from distutils.version import StrictVersion
+from pkg_resources import resource_filename
 from setuptools import (
     Extension,
     find_packages,
@@ -33,43 +34,65 @@ from setuptools import (
 import versioneer
 
 
-class LazyCythonizingList(list):
-    cythonized = False
-
-    def lazy_cythonize(self):
-        if self.cythonized:
-            return
-        self.cythonized = True
-
-        from Cython.Build import cythonize
-        from numpy import get_include
-
-        self[:] = cythonize(
-            [
-                Extension(*ext_args, include_dirs=[get_include()])
-                for ext_args in self
-            ]
+class LazyBuildExtCommandClass(dict):
+    """
+    Lazy command class that defers operations requiring Cython and numpy until
+    they've actually been downloaded and installed by setup_requires.
+    """
+    def __contains__(self, key):
+        return (
+            key == 'build_ext'
+            or super(LazyBuildExtCommandClass, self).__contains__(key)
         )
 
-    def __iter__(self):
-        self.lazy_cythonize()
-        return super(LazyCythonizingList, self).__iter__()
+    def __setitem__(self, key, value):
+        if key == 'build_ext':
+            raise AssertionError("build_ext overridden!")
+        super(LazyBuildExtCommandClass, self).__setitem__(key, value)
 
-    def __getitem__(self, num):
-        self.lazy_cythonize()
-        return super(LazyCythonizingList, self).__getitem__(num)
+    def __getitem__(self, key):
+        if key != 'build_ext':
+            return super(LazyBuildExtCommandClass, self).__getitem__(key)
+
+        from Cython.Distutils import build_ext as cython_build_ext
+
+        class build_ext(cython_build_ext):
+            """
+            Custom build_ext command that lazily adds numpy's include_dir to
+            extensions.
+            """
+            def build_extensions(self):
+                """
+                Lazily append numpy's include directory to Extension includes.
+
+                This is done here rather than at module scope because setup.py
+                may be run before numpy has been installed, in which case
+                importing numpy and calling `numpy.get_include()` will fail.
+                """
+                numpy_incl = resource_filename('numpy', 'core/include')
+                for ext in self.extensions:
+                    ext.include_dirs.append(numpy_incl)
+
+                # This explicitly calls the superclass method rather than the
+                # usual super() invocation because distutils' build_class, of
+                # which Cython's build_ext is a subclass, is an old-style class
+                # in Python 2, which doesn't support `super`.
+                cython_build_ext.build_extensions(self)
+        return build_ext
 
 
-ext_modules = LazyCythonizingList([
-    ('zipline.assets._assets', ['zipline/assets/_assets.pyx']),
-    ('zipline.lib.adjustment', ['zipline/lib/adjustment.pyx']),
-    ('zipline.lib._float64window', ['zipline/lib/_float64window.pyx']),
-    ('zipline.lib._int64window', ['zipline/lib/_int64window.pyx']),
-    ('zipline.lib._uint8window', ['zipline/lib/_uint8window.pyx']),
-    ('zipline.lib.rank', ['zipline/lib/rank.pyx']),
-    ('zipline.data._equities', ['zipline/data/_equities.pyx']),
-    ('zipline.data._adjustments', ['zipline/data/_adjustments.pyx']),
-])
+ext_modules = [
+    Extension('zipline.assets._assets', ['zipline/assets/_assets.pyx']),
+    Extension('zipline.lib.adjustment', ['zipline/lib/adjustment.pyx']),
+    Extension(
+        'zipline.lib._float64window', ['zipline/lib/_float64window.pyx']
+    ),
+    Extension('zipline.lib._int64window', ['zipline/lib/_int64window.pyx']),
+    Extension('zipline.lib._uint8window', ['zipline/lib/_uint8window.pyx']),
+    Extension('zipline.lib.rank', ['zipline/lib/rank.pyx']),
+    Extension('zipline.data._equities', ['zipline/data/_equities.pyx']),
+    Extension('zipline.data._adjustments', ['zipline/data/_adjustments.pyx']),
+]
 
 
 STR_TO_CMP = {
@@ -116,9 +139,8 @@ def _filter_requirements(lines_iter):
             yield requirement
 
 
-REQ_UPPER_BOUNDS = {
-    'numpy': '<1.10',
-}
+# We don't currently have any known upper bounds.
+REQ_UPPER_BOUNDS = {}
 
 
 def _with_bounds(req):
@@ -183,11 +205,12 @@ def extras_requires(conda_format=False):
     }
 
 
-def module_requirements(requirements_path, module_names):
+def module_requirements(requirements_path, module_names, strict_bounds):
     module_names = set(module_names)
     found = set()
     module_lines = []
-    for line in read_requirements(requirements_path, strict_bounds=True):
+    for line in read_requirements(requirements_path,
+                                  strict_bounds=strict_bounds):
         match = REQ_PATTERN.match(line)
         if match is None:
             raise AssertionError("Could not parse requirement: '%s'" % line)
@@ -203,38 +226,12 @@ def module_requirements(requirements_path, module_names):
         )
     return module_lines
 
-
-def pre_setup():
-    if not set(sys.argv) & {'install', 'develop', 'egg_info', 'bdist_wheel'}:
-        return
-
-    try:
-        import pip
-        if StrictVersion(pip.__version__) < StrictVersion('7.1.0'):
-            raise AssertionError(
-                "Zipline installation requires pip>=7.1.0, but your pip "
-                "version is {version}. \n"
-                "You can upgrade your pip with "
-                "'pip install --upgrade pip'.".format(
-                    version=pip.__version__,
-                )
-            )
-    except ImportError:
-        raise AssertionError("Zipline installation requires pip")
-
-    required = ('Cython', 'numpy')
-    for line in module_requirements('etc/requirements.txt', required):
-        pip.main(['install', line])
-
-
-pre_setup()
-
 conda_build = os.path.basename(sys.argv[0]) == 'conda-build'
 
 setup(
     name='zipline',
     version=versioneer.get_version(),
-    cmdclass=versioneer.get_cmdclass(),
+    cmdclass=LazyBuildExtCommandClass(versioneer.get_cmdclass()),
     description='A backtester for financial algorithms.',
     author='Quantopian Inc.',
     author_email='opensource@quantopian.com',
@@ -258,5 +255,10 @@ setup(
     ],
     install_requires=install_requires(conda_format=conda_build),
     extras_require=extras_requires(conda_format=conda_build),
+    setup_requires=module_requirements(
+        'etc/requirements.txt',
+        ('Cython', 'numpy'),
+        strict_bounds=False,
+    ),
     url="http://zipline.io",
 )
diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py
index 8c52986f..459c2538 100644
--- a/tests/pipeline/test_earnings.py
+++ b/tests/pipeline/test_earnings.py
@@ -26,7 +26,7 @@ from zipline.pipeline.loaders.blaze import (
     SID_FIELD_NAME,
     TS_FIELD_NAME,
 )
-from zipline.utils.numpy_utils import make_datetime64D, np_NaT
+from zipline.utils.numpy_utils import make_datetime64D, NaTD
 from zipline.utils.test_utils import (
     make_simple_equity_info,
     tmp_asset_finder,
@@ -234,7 +234,7 @@ class EarningsCalendarLoaderTestCase(TestCase):
 
         # Set NaTs to 0 temporarily because busday_count doesn't support NaT.
         # We fill these entries with NaNs later.
-        whereNaT = raw_announce_dates == np_NaT
+        whereNaT = raw_announce_dates == NaTD
         raw_announce_dates[whereNaT] = make_datetime64D(0)
 
         # The abs call here makes it so that we can use this function to
diff --git a/tests/pipeline/test_engine.py b/tests/pipeline/test_engine.py
index 438b5dbb..29fd2e44 100644
--- a/tests/pipeline/test_engine.py
+++ b/tests/pipeline/test_engine.py
@@ -257,7 +257,7 @@ class ConstantInputTestCase(TestCase):
 
             check_arrays(
                 result['f'].unstack().values,
-                full(result_shape, expected_result),
+                full(result_shape, expected_result, dtype=float),
             )
 
     def test_multiple_rolling_factors(self):
@@ -295,16 +295,16 @@ class ConstantInputTestCase(TestCase):
         # row-wise sum over an array whose values are all (1 - 2)
         check_arrays(
             results['short'].unstack().values,
-            full(shape, -short_factor.window_length),
+            full(shape, -short_factor.window_length, dtype=float),
         )
         check_arrays(
             results['long'].unstack().values,
-            full(shape, -long_factor.window_length),
+            full(shape, -long_factor.window_length, dtype=float),
         )
         # row-wise sum over an array whose values are all (1 - 3)
         check_arrays(
             results['high'].unstack().values,
-            full(shape, -2 * high_factor.window_length),
+            full(shape, -2 * high_factor.window_length, dtype=float),
         )
 
     def test_numeric_factor(self):
@@ -398,13 +398,19 @@ class ConstantInputTestCase(TestCase):
         result_shape = (len(result_index),)
         check_arrays(
             result['sumdiff'],
-            Series(index=result_index, data=full(result_shape, -3)),
+            Series(
+                index=result_index,
+                data=full(result_shape, -3, dtype=float),
+            ),
         )
 
         for name, const in [('open', 1), ('close', 2), ('volume', 3)]:
             check_arrays(
                 result[name],
-                Series(index=result_index, data=full(result_shape, const)),
+                Series(
+                    index=result_index,
+                    data=full(result_shape, const, dtype=float),
+                ),
             )
 
     def test_loader_given_multiple_columns(self):
@@ -471,19 +477,26 @@ class ConstantInputTestCase(TestCase):
                 for name, pipe_col in iteritems(columns)}
 
         index = MultiIndex.from_product([self.dates[2:], self.assets])
+
+        def expected_for_col(col):
+            val = vals[col]
+            offset = columns[col].window_length - min_window
+            return concatenate(
+                [
+                    full(offset * index.levshape[1], nan),
+                    full(
+                        (index.levshape[0] - offset) * index.levshape[1],
+                        val,
+                        float,
+                    )
+                ],
+            )
+
         expected = DataFrame(
-            data={col:
-                  concatenate((
-                      full((columns[col].window_length - min_window)
-                           * index.levshape[1],
-                           nan),
-                      full((index.levshape[0]
-                            - (columns[col].window_length - min_window))
-                           * index.levshape[1],
-                           val)))
-                  for col, val in iteritems(vals)},
+            data={col: expected_for_col(col) for col in vals},
             index=index,
-            columns=columns)
+            columns=columns,
+        )
 
         assert_frame_equal(result, expected)
 
diff --git a/tests/pipeline/test_factor.py b/tests/pipeline/test_factor.py
index 22a2242e..200ed15f 100644
--- a/tests/pipeline/test_factor.py
+++ b/tests/pipeline/test_factor.py
@@ -23,7 +23,12 @@ from zipline.pipeline.factors import (
     RSI,
 )
 from zipline.utils.test_utils import check_allclose, check_arrays
-from zipline.utils.numpy_utils import datetime64ns_dtype, float64_dtype, np_NaT
+from zipline.utils.numpy_utils import (
+    datetime64ns_dtype,
+    float64_dtype,
+    NaTD,
+    NaTns,
+)
 
 from .base import BasePipelineTestCase
 
@@ -309,7 +314,7 @@ class FactorTestCase(BasePipelineTestCase):
         mask = eyemask if use_mask else nomask
         if set_missing:
             asfloat[:, 2] = nan
-            asdatetime[:, 2] = np_NaT
+            asdatetime[:, 2] = NaTns
 
         float_result = masked_rankdata_2d(
             data=asfloat,
@@ -321,7 +326,7 @@ class FactorTestCase(BasePipelineTestCase):
         datetime_result = masked_rankdata_2d(
             data=asdatetime,
             mask=mask,
-            missing_value=np_NaT,
+            missing_value=NaTns,
             method=method,
             ascending=True,
         )
diff --git a/tests/pipeline/test_numerical_expression.py b/tests/pipeline/test_numerical_expression.py
index 1026a762..0315d656 100644
--- a/tests/pipeline/test_numerical_expression.py
+++ b/tests/pipeline/test_numerical_expression.py
@@ -76,9 +76,9 @@ class NumericalExpressionTestCase(TestCase):
         self.h = H()
         self.d = DateFactor()
         self.fake_raw_data = {
-            self.f: full((5, 5), 3),
-            self.g: full((5, 5), 2),
-            self.h: full((5, 5), 1),
+            self.f: full((5, 5), 3, float),
+            self.g: full((5, 5), 2, float),
+            self.h: full((5, 5), 1, float),
             self.d: full((5, 5), 0, dtype='datetime64[ns]'),
         }
         self.mask = DataFrame(True, index=self.dates, columns=self.assets)
@@ -94,7 +94,7 @@ class NumericalExpressionTestCase(TestCase):
 
     def check_constant_output(self, expr, expected):
         self.assertFalse(isnan(expected))
-        return self.check_output(expr, full((5, 5), expected))
+        return self.check_output(expr, full((5, 5), expected, float))
 
     def test_validate_good(self):
         f = self.f
@@ -435,9 +435,9 @@ class NumericalExpressionTestCase(TestCase):
     def test_comparisons(self):
         f, g, h = self.f, self.g, self.h
         self.fake_raw_data = {
-            f: arange(25).reshape(5, 5),
-            g: arange(25).reshape(5, 5) - eye(5),
-            h: full((5, 5), 5),
+            f: arange(25, dtype=float).reshape(5, 5),
+            g: arange(25, dtype=float).reshape(5, 5) - eye(5),
+            h: full((5, 5), 5, dtype=float),
         }
         f_data = self.fake_raw_data[f]
         g_data = self.fake_raw_data[g]
@@ -479,9 +479,9 @@ class NumericalExpressionTestCase(TestCase):
         )
 
         self.fake_raw_data = {
-            f: arange(25).reshape(5, 5),
-            g: arange(25).reshape(5, 5) - eye(5),
-            h: full((5, 5), 5),
+            f: arange(25, dtype=float).reshape(5, 5),
+            g: arange(25, dtype=float).reshape(5, 5) - eye(5),
+            h: full((5, 5), 5, dtype=float),
             custom_filter: custom_filter_mask,
         }
 
diff --git a/tests/pipeline/test_term.py b/tests/pipeline/test_term.py
index f689bc33..c657811a 100644
--- a/tests/pipeline/test_term.py
+++ b/tests/pipeline/test_term.py
@@ -40,6 +40,7 @@ class SomeFactor(Factor):
     dtype = float64_dtype
     window_length = 5
     inputs = [SomeDataSet.foo, SomeDataSet.bar]
+
 SomeFactorAlias = SomeFactor
 
 
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 5c93547b..04dc2ce3 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -16,14 +16,11 @@
 # This code is based on a unittest written by John Salvatier:
 # https://github.com/pymc-devs/pymc/blob/pymc3/tests/test_examples.py
 
-# Disable plotting
-#
-
 import glob
-import imp
 import matplotlib
 from nose_parameterized import parameterized
 import os
+import runpy
 from unittest import TestCase
 
 from zipline.utils import parse_args, run_pipeline
@@ -45,7 +42,7 @@ class ExamplesTests(TestCase):
     @parameterized.expand(((os.path.basename(f).replace('.', '_'), f) for f in
                            glob.glob(os.path.join(example_dir(), '*.py'))))
     def test_example(self, name, example):
-        imp.load_source('__main__', os.path.basename(example), open(example))
+        runpy.run_path(example, run_name='__main__')
 
     # Test algorithm as if scripts/run_algo.py is being used.
     def test_example_run_pipline(self):
diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py
index dcd9b135..6b15f013 100644
--- a/zipline/data/us_equity_pricing.py
+++ b/zipline/data/us_equity_pricing.py
@@ -206,7 +206,7 @@ class BcolzDailyBarWriter(with_metaclass(ABCMeta)):
                 if column_name == 'id':
                     # We know what the content of this column is, so don't
                     # bother reading it.
-                    columns['id'].append(full((nrows,), asset_id))
+                    columns['id'].append(full((nrows,), asset_id, uint32))
                     continue
                 columns[column_name].append(
                     self.to_uint32(table[column_name][:], column_name)
diff --git a/zipline/examples/pairtrade.py b/zipline/examples/pairtrade.py
index d04e8e58..f7fa36c0 100755
--- a/zipline/examples/pairtrade.py
+++ b/zipline/examples/pairtrade.py
@@ -32,8 +32,8 @@ def ols_transform(data, sid1, sid2):
     """Computes regression coefficient (slope and intercept)
     via Ordinary Least Squares between two SIDs.
     """
-    p0 = data.price[sid1]
-    p1 = sm.add_constant(data.price[sid2], prepend=True)
+    p0 = data.price[sid1].values
+    p1 = sm.add_constant(data.price[sid2].values, prepend=True)
     slope, intercept = sm.OLS(p0, p1).fit().params
 
     return slope, intercept
diff --git a/zipline/pipeline/factors/events.py b/zipline/pipeline/factors/events.py
index 42db1fbf..2491efc8 100644
--- a/zipline/pipeline/factors/events.py
+++ b/zipline/pipeline/factors/events.py
@@ -5,7 +5,7 @@ announcements, acquisitions, dividends, etc.).
 from numpy import newaxis
 from zipline.pipeline.data.earnings import EarningsCalendar
 from zipline.utils.numpy_utils import (
-    np_NaT,
+    NaTD,
     busday_count_mask_NaT,
     datetime64D_dtype,
     float64_dtype,
@@ -48,7 +48,7 @@ class BusinessDaysUntilNextEarnings(Factor):
         announce_dates = arrays[0].astype(datetime64D_dtype)
 
         # Set masked values to NaT.
-        announce_dates[~mask] = np_NaT
+        announce_dates[~mask] = NaTD
 
         # Convert row labels into a column vector for broadcasted comparison.
         reference_dates = dates.values.astype(datetime64D_dtype)[:, newaxis]
@@ -84,7 +84,7 @@ class BusinessDaysSincePreviousEarnings(Factor):
         announce_dates = arrays[0].astype(datetime64D_dtype)
 
         # Set masked values to NaT.
-        announce_dates[~mask] = np_NaT
+        announce_dates[~mask] = NaTD
 
         # Convert row labels into a column vector for broadcasted comparison.
         reference_dates = dates.values.astype(datetime64D_dtype)[:, newaxis]
diff --git a/zipline/pipeline/factors/technical.py b/zipline/pipeline/factors/technical.py
index fe08b9c0..8bd59147 100644
--- a/zipline/pipeline/factors/technical.py
+++ b/zipline/pipeline/factors/technical.py
@@ -181,7 +181,7 @@ class _ExponentialWeightedFactor(SingleInputMixin, CustomFactor):
         Return weighting vector for an exponential moving statistic on `length`
         rows with a decay rate of `decay_rate`.
         """
-        return full(length, decay_rate) ** arange(length + 1, 1, -1)
+        return full(length, decay_rate, float) ** arange(length + 1, 1, -1)
 
     @classmethod
     @expect_types(span=Number)
diff --git a/zipline/pipeline/loaders/utils.py b/zipline/pipeline/loaders/utils.py
index e5682461..7f9448e7 100644
--- a/zipline/pipeline/loaders/utils.py
+++ b/zipline/pipeline/loaders/utils.py
@@ -5,7 +5,7 @@ import pandas as pd
 from six import iteritems
 from six.moves import zip
 
-from zipline.utils.numpy_utils import np_NaT
+from zipline.utils.numpy_utils import NaTns
 
 
 def next_date_frame(dates, events_by_sid):
@@ -34,7 +34,7 @@ def next_date_frame(dates, events_by_sid):
     previous_date_frame
     """
     cols = {
-        equity: np.full_like(dates, np_NaT) for equity in events_by_sid
+        equity: np.full_like(dates, NaTns) for equity in events_by_sid
     }
     raw_dates = dates.values
     for equity, event_dates in iteritems(events_by_sid):
@@ -50,7 +50,7 @@ def next_date_frame(dates, events_by_sid):
                 (knowledge_date <= raw_dates) &
                 (raw_dates <= event_date)
             )
-            value_mask = (event_date <= data) | (data == np_NaT)
+            value_mask = (event_date <= data) | (data == NaTns)
             data[date_mask & value_mask] = event_date
 
     return pd.DataFrame(index=dates, data=cols)
@@ -82,7 +82,7 @@ def previous_date_frame(date_index, events_by_sid):
     next_date_frame
     """
     sids = list(events_by_sid)
-    out = np.full((len(date_index), len(sids)), np_NaT, dtype='datetime64[ns]')
+    out = np.full((len(date_index), len(sids)), NaTns, dtype='datetime64[ns]')
     dn = date_index[-1].asm8
     for col_idx, sid in enumerate(sids):
         # events_by_sid[sid] is Series mapping knowledge_date to actual
diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py
index 0c88a625..dfbcde8b 100644
--- a/zipline/utils/factory.py
+++ b/zipline/utils/factory.py
@@ -309,8 +309,7 @@ def create_test_panel_source(sim_params=None, env=None, source_type=None):
                        'arbitrary': arbitrary},
                       index=index)
     if source_type:
-        source_types = np.full(len(index), source_type)
-        df['type'] = source_types
+        df['type'] = source_type
 
     panel = pd.Panel.from_dict({0: df})
 
diff --git a/zipline/utils/numpy_utils.py b/zipline/utils/numpy_utils.py
index 74fed9ae..36b73581 100644
--- a/zipline/utils/numpy_utils.py
+++ b/zipline/utils/numpy_utils.py
@@ -22,11 +22,19 @@ datetime64ns_dtype = dtype('datetime64[ns]')
 
 make_datetime64ns = flip(datetime64, 'ns')
 make_datetime64D = flip(datetime64, 'D')
-np_NaT = make_datetime64ns('NaT')
+
+NaTmap = {
+    dtype('datetime64[%s]' % unit): datetime64('NaT', unit)
+    for unit in ('ns', 'us', 'ms', 's', 'm', 'D')
+}
+NaT_for_dtype = NaTmap.__getitem__
+NaTns = NaT_for_dtype(datetime64ns_dtype)
+NaTD = NaT_for_dtype(datetime64D_dtype)
+
 
 _FILLVALUE_DEFAULTS = {
     float64_dtype: nan,
-    datetime64ns_dtype: np_NaT,
+    datetime64ns_dtype: NaT_for_dtype(datetime64ns_dtype),
 }
 
 
@@ -127,7 +135,10 @@ def repeat_last_axis(array, count):
 _notNaT = make_datetime64D(0)
 
 
-def busday_count_mask_NaT(begindates, enddates, out=None):
+def busday_count_mask_NaT(begindates,
+                          enddates,
+                          out=None,
+                          NaT=NaT_for_dtype(datetime64D_dtype)):
     """
     Simple of numpy.busday_count that returns `float` arrays rather than int
     arrays, and handles `NaT`s by returning `NaN`s where the inputs were `NaT`.
@@ -142,8 +153,8 @@ def busday_count_mask_NaT(begindates, enddates, out=None):
     if out is None:
         out = empty(broadcast(begindates, enddates).shape, dtype=float)
 
-    beginmask = (begindates == np_NaT)
-    endmask = (enddates == np_NaT)
+    beginmask = (begindates == NaT)
+    endmask = (enddates == NaT)
 
     out = busday_count(
         # Temporarily fill in non-NaT values.