Move to .pyi and fix accessing attributes through class (#3)

ikonst · web-flow · commit ba322c139f1a · 2019-01-16T16:22:30.000-08:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@ repos:
       hooks:
           - id: reorder-python-imports
     - repo: https://github.com/pre-commit/pre-commit-hooks
-      rev: v2.0.0
+      rev: v2.1.0
       hooks:
           - id: check-ast
           - id: check-docstring-first
diff --git a/pynamodb_attributes/_typing.pyi b/pynamodb_attributes/_typing.pyi
@@ -0,0 +1,13 @@
+from typing import overload, Any, TypeVar
+
+import pynamodb.attributes
+
+_A = TypeVar('_A', bound=pynamodb.attributes.Attribute)
+_T = TypeVar('_T')
+
+
+class Attribute(pynamodb.attributes.Attribute[_T]):
+    @overload
+    def __get__(self: _A, instance: None, owner: Any) -> _A: ...
+    @overload
+    def __get__(self, instance: Any, owner: Any) -> _T: ...
diff --git a/pynamodb_attributes/integer.py b/pynamodb_attributes/integer.py
@@ -1,6 +1,3 @@
-from typing import Any
-from typing import TYPE_CHECKING
-
 from pynamodb.attributes import Attribute
 from pynamodb.attributes import NumberAttribute
 
@@ -12,7 +9,3 @@ class IntegerAttribute(Attribute):
     attr_type = NumberAttribute.attr_type
     serialize = NumberAttribute.serialize
     deserialize = NumberAttribute.deserialize
-
-    if TYPE_CHECKING:
-        def __get__(self, instance: Any, owner: Any) -> int:
-            ...
diff --git a/pynamodb_attributes/integer.pyi b/pynamodb_attributes/integer.pyi
@@ -0,0 +1,5 @@
+from ._typing import Attribute
+
+
+class IntegerAttribute(Attribute[int]):
+    ...
diff --git a/pynamodb_attributes/integer_date.py b/pynamodb_attributes/integer_date.py
@@ -1,7 +1,5 @@
 import json
 from datetime import date
-from typing import Any
-from typing import TYPE_CHECKING
 
 from pynamodb.attributes import Attribute
 
@@ -18,7 +16,3 @@ def serialize(self, value: date) -> str:
     def deserialize(self, value: str) -> date:
         n = json.loads(value)
         return date(n // 1_00_00, n // 1_00 % 1_00, n % 1_00)
-
-    if TYPE_CHECKING:
-        def __get__(self, instance: Any, owner: Any) -> date:
-            ...
diff --git a/pynamodb_attributes/integer_date.pyi b/pynamodb_attributes/integer_date.pyi
@@ -0,0 +1,7 @@
+from datetime import date
+
+from ._typing import Attribute
+
+
+class IntegerDateAttribute(Attribute[date]):
+    ...
diff --git a/pynamodb_attributes/unicode_delimited_tuple.pyi b/pynamodb_attributes/unicode_delimited_tuple.pyi
@@ -0,0 +1,12 @@
+from typing import Any
+from typing import Type
+from typing import TypeVar
+
+from ._typing import Attribute
+
+_T = TypeVar('_T', bound=tuple)
+
+
+class UnicodeDelimitedTupleAttribute(Attribute[_T]):
+    def __init__(self, tuple_type: Type[_T], delimiter: str = ..., **kwargs: Any) -> None:
+        ...
diff --git a/pynamodb_attributes/unicode_enum.pyi b/pynamodb_attributes/unicode_enum.pyi
@@ -0,0 +1,14 @@
+from enum import Enum
+from typing import Any
+from typing import Optional
+from typing import Type
+from typing import TypeVar
+
+from ._typing import Attribute
+
+_T = TypeVar('_T', bound=Enum)
+
+
+class UnicodeEnumAttribute(Attribute[_T]):
+    def __init__(self, enum_type: Type[_T], unknown_value: Optional[_T] = ..., **kwargs: Any) -> None:
+        ...
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,5 @@
-pre-commit
-pytest
-pytest-cov
-pytest-mock
+mypy>=0.650
+pre-commit>=1.13.0
+pytest>=4.1.1
+pytest-cov>=2.6.1
+pytest-mock>=1.10.0
diff --git a/setup.py b/setup.py
@@ -1,18 +1,31 @@
+import os
+
 from setuptools import find_packages
 from setuptools import setup
 
+
+def find_stubs(package):  # type: ignore
+    stubs = []
+    for root, _, files in os.walk(package):
+        for f in files:
+            path = os.path.join(root, f).replace(package + os.sep, '', 1)
+            if path.endswith('.pyi') or path.endswith('py.typed'):
+                stubs.append(path)
+    return {package: stubs}
+
+
 setup(
     name='pynamodb-attributes',
-    version='0.1.2',
+    version='0.2.0',
     description='Common attributes for PynamoDB',
     url='https://www.github.com/lyft/pynamodb-attributes',
     maintainer='Lyft',
     maintainer_email='ikonstantinov@lyft.com',
     packages=find_packages(exclude=['tests*']),
     dependency_links=[],
     install_requires=[
-        "pynamodb",
+        "pynamodb>=3.3.3",
     ],
     python_requires='>=3',
-    package_data={'pynamodb_attributes': ['py.typed']},
+    package_data=find_stubs('pynamodb_attributes'),
 )
diff --git a/tests/mypy_helpers.py b/tests/mypy_helpers.py
@@ -0,0 +1,49 @@
+import re
+import sys
+from collections import defaultdict
+from tempfile import TemporaryDirectory
+from textwrap import dedent
+from typing import Dict
+from typing import Iterable
+from typing import List
+
+import mypy.api
+
+
+def _run_mypy(program: str) -> Iterable[str]:
+    with TemporaryDirectory() as tempdirname:
+        with open(f'{tempdirname}/__main__.py', 'w') as f:
+            f.write(program)
+        error_pattern = re.compile(fr'^{re.escape(f.name)}:(\d+): error: (.*)$')
+        stdout, stderr, exit_status = mypy.api.run([
+            f.name,
+            '--show-traceback',
+        ])
+        assert exit_status in (0, 1), f"invalid exit status {exit_status}: {stdout}"
+        if stderr:
+            print(stderr, file=sys.stderr)  # allow "printf debugging" of the plugin
+
+        # Group errors by line
+        errors_by_line: Dict[int, List[str]] = defaultdict(list)
+        for line in stdout.split('\n'):
+            m = error_pattern.match(line)
+            if m:
+                errors_by_line[int(m.group(1))].append(m.group(2))
+            elif line:
+                print('x', repr(line))  # allow "printf debugging" of the plugin
+
+        # Reconstruct the "actual" program with "error" comments
+        error_comment_pattern = re.compile(r'(\s+# E: .*)?$')
+        for line_no, line in enumerate(program.split('\n'), start=1):
+            line = error_comment_pattern.sub('', line)
+            errors = errors_by_line.get(line_no)
+            if errors:
+                yield line + ''.join(f'  # E: {error}' for error in errors)
+            else:
+                yield line
+
+
+def assert_mypy_output(program: str) -> None:
+    program = dedent(program).strip()
+    actual = '\n'.join(_run_mypy(program))
+    assert actual == program
diff --git a/tests/mypy_test.py b/tests/mypy_test.py
@@ -0,0 +1,71 @@
+"""
+Note: The expected error strings may change in a future version of mypy.
+      Please update as needed.
+"""
+import pytest
+
+pytest.register_assert_rewrite('tests.mypy_helpers')
+from tests.mypy_helpers import assert_mypy_output  # noqa
+
+
+def test_integer_attribute():
+    assert_mypy_output("""
+    from pynamodb.models import Model
+    from pynamodb_attributes import IntegerAttribute
+
+    class MyModel(Model):
+        my_attr = IntegerAttribute()
+
+    reveal_type(MyModel.my_attr)  # E: Revealed type is 'pynamodb_attributes.integer.IntegerAttribute'
+    reveal_type(MyModel().my_attr)  # E: Revealed type is 'builtins.int*'
+    """)
+
+
+def test_integer_date_attribute():
+    assert_mypy_output("""
+    from pynamodb.models import Model
+    from pynamodb_attributes import IntegerDateAttribute
+
+    class MyModel(Model):
+        my_attr = IntegerDateAttribute()
+
+    reveal_type(MyModel.my_attr)  # E: Revealed type is 'pynamodb_attributes.integer_date.IntegerDateAttribute'
+    reveal_type(MyModel().my_attr)  # E: Revealed type is 'datetime.date*'
+    """)
+
+
+def test_unicode_delimited_tuple_attribute():
+    assert_mypy_output("""
+    from typing import NamedTuple
+    from pynamodb.models import Model
+    from pynamodb_attributes import UnicodeDelimitedTupleAttribute
+
+    class MyTuple(NamedTuple):
+        foo: str
+        bar: str
+
+    class MyModel(Model):
+        my_attr = UnicodeDelimitedTupleAttribute(MyTuple)
+
+    reveal_type(MyModel.my_attr)  # E: Revealed type is 'pynamodb_attributes.unicode_delimited_tuple.UnicodeDelimitedTupleAttribute[Tuple[builtins.str, builtins.str, fallback=__main__.MyTuple]]'
+    reveal_type(MyModel().my_attr)  # E: Revealed type is 'Tuple[builtins.str, builtins.str, fallback=__main__.MyTuple]'
+    reveal_type(MyModel().my_attr.foo)  # E: Revealed type is 'builtins.str'
+    """)  # noqa: E501
+
+
+def test_unicode_enum_attribute():
+    assert_mypy_output("""
+    from enum import Enum
+    from pynamodb.models import Model
+    from pynamodb_attributes import UnicodeEnumAttribute
+
+    class MyEnum(Enum):
+        foo = 'foo'
+        bar = 'bar'
+
+    class MyModel(Model):
+        my_attr = UnicodeEnumAttribute(MyEnum)
+
+    reveal_type(MyModel.my_attr)  # E: Revealed type is 'pynamodb_attributes.unicode_enum.UnicodeEnumAttribute[__main__.MyEnum]'
+    reveal_type(MyModel().my_attr)  # E: Revealed type is '__main__.MyEnum*'
+    """)  # noqa: E501

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +from ._typing import Attribute
++
++
 +class IntegerAttribute(Attribute[int]):
 +    ...