extmod/modtyping: Add support for tytping aliases and typing.Literal …

…in typing alias specifications.

Signed-off-by: Jos Verlinde <[email protected]>

extmod/modtyping.c

@@ -25,12 +25,14 @@
  */
 
 #include "py/obj.h"
+#include "py/runtime.h"
 
 #if MICROPY_PY_TYPING
 
 // Implement roughly the equivalent of the following minimal Python typing module, meant to support the
 // typing syntax at runtime but otherwise ignoring any functionality:
 //
+//  TYPE_CHECKING = False
 //  class _AnyCall:
 //    def __init__(*args, **kwargs):
 //        pass
@@ -59,6 +61,78 @@ typedef struct _mp_obj_any_call_t
 } mp_obj_any_call_t;
 
 static const mp_obj_type_t mp_type_any_call_t;
+static const mp_obj_type_t mp_type_typing_alias;
+
+// Lightweight runtime representation for objects such as typing.List[int].
+// The alias keeps track of the original builtin type and the tuple of
+// parameters so that __origin__ and __args__ can be queried at runtime.
+typedef struct _mp_obj_typing_alias_t {
+    mp_obj_base_t base;
+    mp_obj_t origin;
+    mp_obj_t args; // tuple or MP_OBJ_NULL when not parametrised
+} mp_obj_typing_alias_t;
+
+// Maps a qstr name to the builtin type that should back the alias.
+typedef struct {
+    qstr name;
+    const mp_obj_type_t *type;
+} typing_alias_spec_t;
+
+static mp_obj_t typing_alias_from_spec(const typing_alias_spec_t *spec_table, size_t spec_len, qstr attr);
+
+static mp_obj_t typing_alias_new(mp_obj_t origin, mp_obj_t args) {
+    mp_obj_typing_alias_t *self = mp_obj_malloc(mp_obj_typing_alias_t, &mp_type_typing_alias);
+    self->origin = origin;
+    self->args = args;
+    return MP_OBJ_FROM_PTR(self);
+}
+
+static void typing_alias_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
+    // Only handle reads that we recognise: __origin__ and __args__.  Anything
+    // else is delegated back to the VM where it will fall through to the
+    // generic AnyCall behaviour.
+    if (dest[0] != MP_OBJ_NULL) {
+        return;
+    }
+
+    mp_obj_typing_alias_t *self = MP_OBJ_TO_PTR(self_in);
+    if (attr == MP_QSTR___origin__) {
+        dest[0] = self->origin;
+    } else if (attr == MP_QSTR___args__) {
+        dest[0] = self->args == MP_OBJ_NULL ? mp_const_empty_tuple : self->args;
+    }
+}
+
+static mp_obj_t typing_alias_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value) {
+    if (value != MP_OBJ_SENTINEL) {
+        mp_raise_TypeError(MP_ERROR_TEXT("typing alias does not support assignment"));
+    }
+
+    mp_obj_typing_alias_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_obj_t args_obj;
+    if (mp_obj_is_type(index_in, &mp_type_tuple)) {
+        args_obj = index_in;
+    } else {
+        mp_obj_t items[1] = { index_in };
+        args_obj = mp_obj_new_tuple(1, items);
+    }
+
+    return typing_alias_new(self->origin, args_obj);
+}
+
+static mp_obj_t typing_alias_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
+    mp_obj_typing_alias_t *self = MP_OBJ_TO_PTR(self_in);
+    return mp_call_function_n_kw(self->origin, n_args, n_kw, args);
+}
+
+static MP_DEFINE_CONST_OBJ_TYPE(
+    mp_type_typing_alias,
+    MP_QSTR_typing_alias,
+    MP_TYPE_FLAG_NONE,
+    attr, typing_alias_attr,
+    subscr, typing_alias_subscr,
+    call, typing_alias_call
+    );
 
 
 // Can be used both for __new__ and __call__: the latter's prototype is
@@ -114,10 +188,39 @@ static void any_call_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     }
 }
 
+// Only a small subset of typing.* names need concrete runtime behaviour.  The
+// table below lists those names together with the builtin type that should be
+// wrapped in a typing alias.  Everything else continues to use the extremely
+// small AnyCall shim.
+static const typing_alias_spec_t typing_container_specs[] = {
+    { MP_QSTR_type, &mp_type_type },
+    { MP_QSTR_Type, &mp_type_type },
+    { MP_QSTR_List, &mp_type_list },
+    { MP_QSTR_Dict, &mp_type_dict },
+    { MP_QSTR_Tuple, &mp_type_tuple },
+    { MP_QSTR_Literal, &mp_type_any_call_t },
+    #if MICROPY_PY_BUILTINS_SET
+    { MP_QSTR_Set, &mp_type_set },
+    #endif
+    #if MICROPY_PY_BUILTINS_FROZENSET
+    { MP_QSTR_FrozenSet, &mp_type_frozenset },
+    #endif
+};
+
 void any_call_module_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     // Only loading is supported.
     if (dest[0] == MP_OBJ_NULL) {
-        dest[0] = MP_OBJ_FROM_PTR(&mp_type_any_call_t);
+        // First see if this attribute corresponds to a container alias that
+        // needs a proper __getitem__ implementation.
+        mp_obj_t alias = typing_alias_from_spec(typing_container_specs, MP_ARRAY_SIZE(typing_container_specs), attr);
+        if (alias != MP_OBJ_NULL) {
+            dest[0] = alias;
+        } else {
+            // Otherwise fall back to returning the singleton AnyCall object,
+            // preserving the "typing ignores everything" behaviour used for
+            // the majority of names.
+            dest[0] = MP_OBJ_FROM_PTR(&mp_type_any_call_t);
+        }
     }
 }
 
@@ -131,9 +234,21 @@ static MP_DEFINE_CONST_OBJ_TYPE(
     call, any_call_call
     );
 
+// Helper to look up a qstr in the alias specification table and lazily create
+// the corresponding typing alias object when a match is found.
+static mp_obj_t typing_alias_from_spec(const typing_alias_spec_t *spec_table, size_t spec_len, qstr attr) {
+    for (size_t i = 0; i < spec_len; ++i) {
+        if (spec_table[i].name == attr) {
+            mp_obj_t origin = MP_OBJ_FROM_PTR(spec_table[i].type);
+            return typing_alias_new(origin, MP_OBJ_NULL);
+        }
+    }
+    return MP_OBJ_NULL;
+}
 
 static const mp_rom_map_elem_t mp_module_typing_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_typing) },
+    { MP_ROM_QSTR(MP_QSTR_TYPE_CHECKING), MP_ROM_FALSE },
 };
 
 static MP_DEFINE_CONST_DICT(mp_module_typing_globals, mp_module_typing_globals_table);