Added typing information for datetime64, timedelta64 (#31)

Author: Matthew Bedder

numpy-stubs/__init__.pyi

@@ -1,5 +1,6 @@
 import builtins
 import sys
+import datetime as dt
 
 from numpy.core._internal import _ctypes
 from typing import (
@@ -393,14 +394,42 @@ class number(generic):
 
 class bool_(_real_generic): ...
 class object_(generic): ...
-class datetime64(_real_generic): ...
+
+class datetime64:
+    @overload
+    def __init__(
+        self, _data: Union[datetime64, str, dt.datetime] = ..., _format: str = ...
+    ) -> None: ...
+    @overload
+    def __init__(self, _data: int, _format: str) -> None: ...
+    def __add__(self, other: Union[timedelta64, int]) -> datetime64: ...
+    def __sub__(self, other: Union[timedelta64, datetime64, int]) -> timedelta64: ...
+
 class integer(number, _real_generic): ...
 class signedinteger(integer): ...
 class int8(signedinteger): ...
 class int16(signedinteger): ...
 class int32(signedinteger): ...
 class int64(signedinteger): ...
-class timedelta64(signedinteger): ...
+
+class timedelta64(signedinteger):
+    def __init__(self, _data: Any = ..., _format: str = ...) -> None: ...
+    @overload
+    def __add__(self, other: Union[timedelta64, int]) -> timedelta64: ...
+    @overload
+    def __add__(self, other: datetime64) -> datetime64: ...
+    def __sub__(self, other: Union[timedelta64, int]) -> timedelta64: ...
+    if sys.version_info[0] < 3:
+        @overload
+        def __div__(self, other: timedelta64) -> float: ...
+        @overload
+        def __div__(self, other: float) -> timedelta64: ...
+    @overload
+    def __truediv__(self, other: timedelta64) -> float: ...
+    @overload
+    def __truediv__(self, other: float) -> timedelta64: ...
+    def __mod__(self, other: timedelta64) -> timedelta64: ...
+
 class unsignedinteger(integer): ...
 class uint8(unsignedinteger): ...
 class uint16(unsignedinteger): ...

tests/fail/scalars.py

@@ -7,3 +7,15 @@
 
 np.complex64(1, 2)  # E: Too many arguments
 # TODO: protocols (can't check for non-existent protocols w/ __getattr__)
+
+np.datetime64(0)  # E: non-matching overload
+
+dt_64 = np.datetime64(0, 'D')
+td_64 = np.timedelta64(1, 'h')
+
+dt_64 + dt_64  # E: Unsupported operand types
+
+td_64 - dt_64  # E: Unsupported operand types
+td_64 / dt_64  # E: No overload
+td_64 % 1  # E: Unsupported operand types
+td_64 % dt_64  # E: Unsupported operand types

tests/pass/scalars.py

@@ -50,3 +50,26 @@ def __float__(self):
 np.uint16().ndim + 1
 np.uint32().strides
 np.uint64().shape
+
+# Time structures
+np.datetime64()
+np.datetime64(0, 'D')
+np.datetime64('2019')
+np.datetime64('2019', 'D')
+
+np.timedelta64()
+np.timedelta64(0)
+np.timedelta64(0, 'D')
+
+dt_64 = np.datetime64(0, 'D')
+td_64 = np.timedelta64(1, 'h')
+
+dt_64 + td_64
+dt_64 - dt_64
+dt_64 - td_64
+
+td_64 + td_64
+td_64 - td_64
+td_64 / 1.
+td_64 / td_64
+td_64 % td_64

tests/reveal/scalars.py

@@ -11,3 +11,20 @@
 reveal_type(x.itemsize)   # E: int
 reveal_type(x.shape)      # E: tuple[builtins.int]
 reveal_type(x.strides)    # E: tuple[builtins.int]
+
+# Time structures
+dt = np.datetime64(0, 'D')
+td = np.timedelta64(0, 'D')
+
+reveal_type(dt + td)  # E: numpy.datetime64
+reveal_type(dt + 1)   # E: numpy.datetime64
+reveal_type(dt - dt)  # E: numpy.timedelta64
+reveal_type(dt - 1)   # E: numpy.timedelta64
+
+reveal_type(td + td)  # E: numpy.timedelta64
+reveal_type(td + 1)   # E: numpy.timedelta64
+reveal_type(td - td)  # E: numpy.timedelta64
+reveal_type(td - 1)   # E: numpy.timedelta64
+reveal_type(td / 1.)  # E: numpy.timedelta64
+reveal_type(td / td)  # E: float
+reveal_type(td % td)  # E: numpy.timedelta64

tests/test_stubs.py

@@ -1,4 +1,5 @@
 import os
+from collections import defaultdict
 
 import pytest
 from mypy import api
@@ -50,14 +51,14 @@ def test_fail(path, py2_arg):
     with open(path) as fin:
         lines = fin.readlines()
 
-    errors = {}
+    errors = defaultdict(lambda: "")
     for error_line in stdout.split("\n"):
         error_line = error_line.strip()
         if not error_line:
             continue
 
         lineno = int(error_line.split(":")[1])
-        errors[lineno] = error_line
+        errors[lineno] += error_line
 
     for i, line in enumerate(lines):
         lineno = i + 1