From d0a31458a85a436fb0789d013e188ffec4ae4b14 Mon Sep 17 00:00:00 2001
From: Tim Rid <6593626+timrid@users.noreply.github.com>
Date: Sat, 13 Mar 2021 17:49:03 +0100
Subject: [PATCH] added covariant and contravariant from PR #1

---
 construct-stubs/core.pyi           | 12 +++++------
 construct_typed/generic_wrapper.py |  8 +++----
 construct_typed/tstruct.py         |  2 +-
 tests/test_core.py                 | 34 ++++++++++++++++++++++++++++++
 4 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/construct-stubs/core.pyi b/construct-stubs/core.pyi
index 95d90bf..2214e02 100644
--- a/construct-stubs/core.pyi
+++ b/construct-stubs/core.pyi
@@ -84,8 +84,8 @@ def stream_iseof(stream: t.BinaryIO) -> bool: ...
 # ===============================================================================
 # abstract constructs
 # ===============================================================================
-ParsedType = t.TypeVar("ParsedType")
-BuildTypes = t.TypeVar("BuildTypes")
+ParsedType = t.TypeVar("ParsedType", covariant=True)
+BuildTypes = t.TypeVar("BuildTypes", contravariant=True)
 
 class Construct(t.Generic[ParsedType, BuildTypes]):
     name: t.Optional[str]
@@ -151,8 +151,8 @@ ValueType = t.TypeVar("ValueType")
 ConstantOrContextLambda = t.Union[ValueType, t.Callable[[Context], t.Any]]
 ConstantOrContextLambda2 = t.Union[ValueType, t.Callable[[Context], ValueType]]
 
-SubconParsedType = t.TypeVar("SubconParsedType")
-SubconBuildTypes = t.TypeVar("SubconBuildTypes")
+SubconParsedType = t.TypeVar("SubconParsedType", covariant=True)
+SubconBuildTypes = t.TypeVar("SubconBuildTypes", contravariant=True)
 
 class Subconstruct(
     t.Generic[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes],
@@ -547,11 +547,11 @@ class Rebuild(Subconstruct[SubconParsedType, SubconBuildTypes, ParsedType, Build
     ) -> Rebuild[SubconParsedType, SubconBuildTypes, SubconParsedType, None]: ...
 
 class Default(Subconstruct[SubconParsedType, SubconBuildTypes, ParsedType, BuildTypes]):
-    value: ConstantOrContextLambda2[SubconBuildTypes]
+    value: ConstantOrContextLambda[SubconBuildTypes]
     def __new__(
         cls,
         subcon: Construct[SubconParsedType, SubconBuildTypes],
-        value: ConstantOrContextLambda2[SubconBuildTypes],
+        value: ConstantOrContextLambda[SubconBuildTypes],
     ) -> Default[
         SubconParsedType,
         SubconBuildTypes,
diff --git a/construct_typed/generic_wrapper.py b/construct_typed/generic_wrapper.py
index 197657e..d3dab33 100644
--- a/construct_typed/generic_wrapper.py
+++ b/construct_typed/generic_wrapper.py
@@ -3,10 +3,10 @@ import enum
 import textwrap
 import typing as t
 
-ParsedType = t.TypeVar("ParsedType")
-BuildTypes = t.TypeVar("BuildTypes")
-SubconParsedType = t.TypeVar("SubconParsedType")
-SubconBuildTypes = t.TypeVar("SubconBuildTypes")
+ParsedType = t.TypeVar("ParsedType", covariant=True)
+BuildTypes = t.TypeVar("BuildTypes", contravariant=True)
+SubconParsedType = t.TypeVar("SubconParsedType", covariant=True)
+SubconBuildTypes = t.TypeVar("SubconBuildTypes", contravariant=True)
 ListType = t.TypeVar("ListType")
 ValueType = t.TypeVar("ValueType")
 
diff --git a/construct_typed/tstruct.py b/construct_typed/tstruct.py
index 66c3b08..fab73e1 100644
--- a/construct_typed/tstruct.py
+++ b/construct_typed/tstruct.py
@@ -51,7 +51,7 @@ class TContainerBase(_TContainerBase):
 
 
 def TStructField(
-    subcon: Construct[ParsedType, BuildTypes],
+    subcon: Construct[ParsedType, t.Any],
     doc: t.Optional[str] = None,
     parsed: t.Optional[t.Callable[[t.Any, Context], None]] = None,
 ) -> ParsedType:
diff --git a/tests/test_core.py b/tests/test_core.py
index cd5671f..78c6740 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -2227,3 +2227,37 @@ def test_struct_issue_771() -> None:
     assert spec.sizeof(**info) == 10
 
 
+def test_buildtypes_contravariance() -> None:
+    if t.TYPE_CHECKING:
+        class HexString(Adapter[bytes, bytes, str, str]): ...
+    else:
+        class HexString(Adapter):
+            def _decode(self, obj, context, path):
+                return obj.hex()
+
+            def _encode(self, obj, context, path):
+                return bytes.fromhex(obj)
+
+    HexStringBytes = HexString(Bytes(2))
+    assert HexStringBytes.build('1234') == b'\x12\x34'
+    assert HexStringBytes.parse(b'\x56\x78') == '5678'
+
+    # this fails if BuildTypes is not contravariant,
+    # as GreedyBytes has a BuildType of Union[bytes, int]
+    HexStringGreedyBytes = HexString(GreedyBytes)
+    assert HexStringGreedyBytes.build('9abc') == b'\x9a\xbc'
+    assert HexStringGreedyBytes.parse(b'\xcd\xef') == 'cdef'
+
+
+def test_parsetype_covariance() -> None:
+    T = t.TypeVar('T')
+    if t.TYPE_CHECKING:
+        class ReversedList(SymmetricAdapter[t.List[T], t.List[T], t.List[T], t.List[T]]): ...
+    else:
+        class ReversedList(SymmetricAdapter):
+            def _decode(self, obj, context, path):
+                return list(reversed(obj))
+
+    assert ReversedList(Array(4, Byte)).build([1, 2, 3, 4]) == b'\x04\x03\x02\x01'
+    assert ReversedList(Array(4, Byte)).parse(b'\x01\x02\x03\x04') == [4, 3, 2, 1]
+    
\ No newline at end of file