Repository: polygraph-python/polygraph

Issue:
Use square brackets for wrapping types, e.g. `Union[X, Y]`

Comments:

Patch:
diff --git a/polygraph/types/type_builder.py b/polygraph/types/type_builder.py
index 2b20577..fcce7f3 100644
--- a/polygraph/types/type_builder.py
+++ b/polygraph/types/type_builder.py
@@ -2,7 +2,11 @@ from functools import wraps
 
 from polygraph.exceptions import PolygraphSchemaError, PolygraphValueError
 from polygraph.types.api import typedef
-from polygraph.types.basic_type import PolygraphOutputType, PolygraphType
+from polygraph.types.basic_type import (
+    PolygraphOutputType,
+    PolygraphType,
+    PolygraphTypeMeta,
+)
 from polygraph.types.definitions import TypeKind
 from polygraph.utils.deduplicate import deduplicate
 
@@ -39,7 +43,15 @@ def type_builder_cache(method):
     return wrapper
 
 
-class Union(PolygraphOutputType, PolygraphType):
+class TypeBuilderMeta(PolygraphTypeMeta):
+    def __getitem__(self, value):
+        try:
+            return self.__new__(self, *value)
+        except TypeError:
+            return self.__new__(self, value)
+
+
+class Union(PolygraphOutputType, PolygraphType, metaclass=TypeBuilderMeta):
     """
     GraphQL Unions represent an object that could be one of a list of
     GraphQL Object types, but provides for no guaranteed fields between
@@ -77,7 +89,7 @@ class Union(PolygraphOutputType, PolygraphType):
         return type(name, bases, attrs)
 
 
-class List(PolygraphOutputType, PolygraphType):
+class List(PolygraphOutputType, PolygraphType, metaclass=TypeBuilderMeta):
     """
     A GraphQL list is a special collection type which declares the
     type of each item in the List (referred to as the item type of
@@ -109,7 +121,7 @@ class List(PolygraphOutputType, PolygraphType):
         return type(name, bases, attrs)
 
 
-class NonNull(PolygraphType):
+class NonNull(PolygraphType, metaclass=TypeBuilderMeta):
     """
     Represents a type for which null is not a valid result.
     """


Test patch:
diff --git a/polygraph/types/tests/helper.py b/polygraph/types/tests/helper.py
new file mode 100644
index 0000000..59d17f2
--- /dev/null
+++ b/polygraph/types/tests/helper.py
@@ -0,0 +1,23 @@
+from polygraph.types.decorators import field
+from polygraph.types.object_type import ObjectType
+from polygraph.types.scalar import Int, String
+
+
+class Person(ObjectType):
+    @field()
+    def name(self) -> String:
+        pass
+
+    @field()
+    def age(self) -> Int:
+        pass
+
+
+class Animal(ObjectType):
+    @field()
+    def name(self) -> String:
+        pass
+
+    @field()
+    def sound(self) -> String:
+        pass
diff --git a/polygraph/types/tests/test_types.py b/polygraph/types/tests/test_types.py
index d08d30f..cb6bf0a 100644
--- a/polygraph/types/tests/test_types.py
+++ b/polygraph/types/tests/test_types.py
@@ -50,6 +50,9 @@ class NonNullTest(TestCase):
         with self.assertRaises(TypeError):
             NonNull(NonNullString)
 
+    def test_square_bracket_notation(self):
+        self.assertEqual(NonNull(String), NonNull[String])
+
 
 class ListTest(TestCase):
 
@@ -67,3 +70,6 @@ class ListTest(TestCase):
         self.assertEqual(string_list(["a", "b", "c"]), ["a", "b", "c"])
         with self.assertRaises(PolygraphValueError):
             string_list(["a", "b", "c", None])
+
+    def test_square_bracket_notation(self):
+        self.assertEqual(List(Int), List[Int])
diff --git a/polygraph/types/tests/test_union.py b/polygraph/types/tests/test_union.py
index e98ba92..1150486 100644
--- a/polygraph/types/tests/test_union.py
+++ b/polygraph/types/tests/test_union.py
@@ -1,17 +1,23 @@
-from unittest import TestCase, skip
+from unittest import TestCase
 
 from polygraph.exceptions import PolygraphValueError
 from polygraph.types.scalar import Float, Int, String
+from polygraph.types.tests.helper import Animal, Person
 from polygraph.types.type_builder import Union
 
 
 class UnionTypeTest(TestCase):
 
+    def test_square_bracket_notation(self):
+        self.assertEqual(
+            Union(Person, Animal),
+            Union[Person, Animal],
+        )
+
     def test_commutativity(self):
         self.assertEqual(Union(String, Int), Union(Int, String))
         self.assertEqual(Union(String, Int, Float), Union(Float, String, Int))
 
-    @skip
     def test_associativity(self):
         self.assertEqual(
             Union(Union(String, Int), Float),