From 76d7cbf36ee0c1c89dfd26ae64ded6c6a0bc2e36 Mon Sep 17 00:00:00 2001
From: Daniel Rosenwasser <Daniel.Rosenwasser@microsoft.com>
Date: Mon, 8 Jan 2024 13:44:22 -0800
Subject: [PATCH] Support abstract collection types.

---
 .../ts_conversion/python_type_to_ts_nodes.py  | 34 +++++++++----------
 python/tests/abstract_collections.py          | 29 ++++++++++++++++
 2 files changed, 46 insertions(+), 17 deletions(-)
 create mode 100644 python/tests/abstract_collections.py

diff --git a/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py b/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py
index c41425c5..f66d4b5f 100644
--- a/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py
+++ b/python/src/typechat/_internal/ts_conversion/python_type_to_ts_nodes.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import typing
+import collections.abc
 from dataclasses import dataclass
 from types import NoneType, UnionType, get_original_bases
 from typing import (
@@ -16,6 +16,7 @@
     NotRequired,
     Protocol,
     Required,
+    Self,
     TypeAlias,
     TypeAliasType,
     TypedDict,
@@ -94,25 +95,24 @@ class TypeScriptNodeTranslationResult:
     errors: list[str]
 
 
-# TODO: https://github.com/microsoft/pyright/issues/6587
-_SELF_TYPE = getattr(typing, "Self")
-
 _LIST_TYPES: set[object] = {
     list,
     set,
     frozenset,
-    # TODO: https://github.com/microsoft/pyright/issues/6582
-    # collections.abc.MutableSequence,
-    # collections.abc.Sequence,
-    # collections.abc.Set
+    collections.abc.Sequence,
+    collections.abc.MutableSequence,
+    collections.abc.Set,
+    collections.abc.MutableSet,
+    collections.abc.Iterable,
+    collections.abc.Collection,
 }
 
-# TODO: https://github.com/microsoft/pyright/issues/6582
-# _DICT_TYPES: set[type] = {
-#     dict,
-#     collections.abc.MutableMapping,
-#     collections.abc.Mapping
-# }
+
+_DICT_TYPES: set[type] = {
+    dict,
+    collections.abc.MutableMapping,
+    collections.abc.Mapping,
+}
 
 
 def python_type_to_typescript_nodes(root_py_type: object) -> TypeScriptNodeTranslationResult:
@@ -184,12 +184,12 @@ def convert_to_type_node(py_type: object) -> TypeNode:
             return NullTypeReferenceNode
         if py_type is Never or py_type is NoReturn:
             return NeverTypeReferenceNode
-        if py_type is _SELF_TYPE:
+        if py_type is Self:
             return ThisTypeReferenceNode
 
         # TODO: consider handling bare 'tuple' (and list, etc.)
         # https://docs.python.org/3/library/typing.html#annotating-tuples
-        # Using plain tuple as an annotation is equivalent to using tuple[Any, ...]:
+        # Using plain 'tuple' as an annotation is equivalent to using 'tuple[Any, ...]':
 
         origin = get_origin(py_type)
         if origin is not None:
@@ -199,7 +199,7 @@ def convert_to_type_node(py_type: object) -> TypeNode:
                     return TypeReferenceNode(IdentifierNode("Array"), [type_arg])
                 return ArrayTypeNode(type_arg)
 
-            if origin is dict:
+            if origin in _DICT_TYPES:
                 # TODO
                 # Currently, we naively assume all dicts are string-keyed
                 # unless they're annotated with `int` or `float` (note: not `int | float`).
diff --git a/python/tests/abstract_collections.py b/python/tests/abstract_collections.py
new file mode 100644
index 00000000..1435f8c0
--- /dev/null
+++ b/python/tests/abstract_collections.py
@@ -0,0 +1,29 @@
+
+from typing import TypedDict
+from typechat import python_type_to_typescript_schema
+
+import collections.abc
+
+class MyType(TypedDict):
+    built_in_dict: dict[str, str]
+    built_in_set: set[str]
+    built_in_frozen_set: frozenset[str]
+
+    mapping: collections.abc.Mapping[str, str]
+    mutable_mapping: collections.abc.MutableMapping[str, str]
+
+    set: collections.abc.Set[str]
+    mutable_set: collections.abc.MutableSet[str]
+
+    sequence: collections.abc.Sequence[str]
+    mutable_sequence: collections.abc.MutableSequence[str]
+
+result = python_type_to_typescript_schema(MyType)
+
+print(f"// Entry point is: '{result.typescript_type_reference}'")
+print("// TypeScript Schema:\n")
+print(result.typescript_schema_str)
+if result.errors:
+    print("// Errors:")
+    for err in result.errors:
+        print(f"// - {err}\n")