feat(core): Add common methods to list/dict (#9)

potatomashed · web-flow · commit 02affe557295 · 2025-01-20T09:09:03.000-08:00
diff --git a/.gitignore b/.gitignore
@@ -2,3 +2,4 @@
 *.dSYM
 build
 build-cpp-tests
+python/mlc/_version.py
diff --git a/cpp/c_api_tests.cc b/cpp/c_api_tests.cc
@@ -311,9 +311,8 @@ struct FieldSetter {
 
 MLC_REGISTER_FUNC("mlc.testing.FieldGet").set_body([](ObjectRef root, const char *target_name) {
   Any ret;
-  MLCTypeInfo *info = ::mlc::Lib::GetTypeInfo(root.GetTypeIndex());
   try {
-    ::mlc::core::VisitFields(root.get(), info, FieldGetter{target_name, &ret});
+    ::mlc::core::VisitFields(root.get(), nullptr, FieldGetter{target_name, &ret});
   } catch (FieldFoundException &) {
     return ret;
   }
@@ -322,9 +321,8 @@ MLC_REGISTER_FUNC("mlc.testing.FieldGet").set_body([](ObjectRef root, const char
 });
 
 MLC_REGISTER_FUNC("mlc.testing.FieldSet").set_body([](ObjectRef root, const char *target_name, Any src) {
-  MLCTypeInfo *info = ::mlc::Lib::GetTypeInfo(root.GetTypeIndex());
   try {
-    ::mlc::core::VisitFields(root.get(), info, FieldSetter{target_name, src});
+    ::mlc::core::VisitFields(root.get(), nullptr, FieldSetter{target_name, src});
   } catch (FieldFoundException &) {
     return;
   }
diff --git a/include/mlc/core/all.h b/include/mlc/core/all.h
@@ -1,17 +1,17 @@
 #ifndef MLC_CORE_ALL_H_
 #define MLC_CORE_ALL_H_
-#include "./dict.h"          // IWYU pragma: export
-#include "./error.h"         // IWYU pragma: export
-#include "./field_visitor.h" // IWYU pragma: export
-#include "./func.h"          // IWYU pragma: export
-#include "./func_details.h"  // IWYU pragma: export
-#include "./list.h"          // IWYU pragma: export
-#include "./object.h"        // IWYU pragma: export
-#include "./object_path.h"   // IWYU pragma: export
-#include "./reflection.h"    // IWYU pragma: export
-#include "./str.h"           // IWYU pragma: export
-#include "./typing.h"        // IWYU pragma: export
-#include "./utils.h"         // IWYU pragma: export
+#include "./dict.h"         // IWYU pragma: export
+#include "./error.h"        // IWYU pragma: export
+#include "./func.h"         // IWYU pragma: export
+#include "./func_details.h" // IWYU pragma: export
+#include "./list.h"         // IWYU pragma: export
+#include "./object.h"       // IWYU pragma: export
+#include "./object_path.h"  // IWYU pragma: export
+#include "./reflection.h"   // IWYU pragma: export
+#include "./str.h"          // IWYU pragma: export
+#include "./typing.h"       // IWYU pragma: export
+#include "./utils.h"        // IWYU pragma: export
+#include "./visitor.h"      // IWYU pragma: export
 
 namespace mlc {
 namespace core {
diff --git a/include/mlc/core/dict.h b/include/mlc/core/dict.h
@@ -73,8 +73,8 @@ struct UDictObj : protected ::mlc::core::DictBase {
   MLC_INLINE iterator find(const Any &key) { return iterator(IterStateMut{this, Acc::Find(this, key)}); }
   MLC_INLINE const_iterator find(const Any &key) const { return const_iterator(IterStateConst{this, Acc::Find(this, key)}); }
   MLC_INLINE void erase(const Any &key) { Acc::Erase(this, key); }
-  MLC_INLINE void erase(const iterator &it) { Acc::Erase(this, it.i.i); }
-  MLC_INLINE void erase(const const_iterator &it) { Acc::Erase(this, it.i.i); }
+  MLC_INLINE void erase(const iterator &it) { Acc::_Erase(this, it.i.i); }
+  MLC_INLINE void erase(const const_iterator &it) { Acc::_Erase(this, it.i.i); }
   // clang-format on
   template <typename K, typename V> MLC_INLINE_NO_MSVC DictObj<K, V> *AsTyped() const;
 
@@ -149,6 +149,8 @@ struct UDict : public ObjectRef {
       .StaticFn("__init__", FromAnyTuple)
       .MemFn("__str__", &UDictObj::__str__)
       .MemFn("__getitem__", ::mlc::core::DictBase::Accessor<UDictObj>::GetItem)
+      .MemFn("__setitem__", ::mlc::core::DictBase::Accessor<UDictObj>::SetItem)
+      .MemFn("__delitem__", ::mlc::core::DictBase::Accessor<UDictObj>::Erase)
       .MemFn("__iter_get_key__", ::mlc::core::DictBase::Accessor<UDictObj>::GetKey)
       .MemFn("__iter_get_value__", ::mlc::core::DictBase::Accessor<UDictObj>::GetValue)
       .MemFn("__iter_advance__", ::mlc::core::DictBase::Accessor<UDictObj>::Advance);
diff --git a/include/mlc/core/dict_base.h b/include/mlc/core/dict_base.h
@@ -39,8 +39,8 @@ struct DictBase : public MLCDict {
     inline static void WithCapacity(TDictObj *self, int64_t new_cap);
     inline static KVPair *InsertOrLookup(TDictObj *self, Any key);
     inline static KVPair *TryInsertOrLookup(TDictObj *self, MLCAny *key);
-    inline static void Erase(TDictObj *self, const Any &key);
-    inline static void Erase(TDictObj *self, int64_t index);
+    inline static Any Erase(TDictObj *self, const Any &key);
+    inline static void _Erase(TDictObj *self, int64_t index);
     inline static Any &At(TDictObj *self, const Any &key);
     inline static const Any &At(const TDictObj *self, const Any &key);
     inline static Any &Bracket(TDictObj *self, const Any &key) {
@@ -52,6 +52,8 @@ struct DictBase : public MLCDict {
     inline static BlockIter Prev(const TDictObj *self, BlockIter iter);
     inline static void New(int32_t num_args, const AnyView *args, Any *any_ret);
     inline static Any GetItem(TDictObj *self, Any key) { return self->at(key); }
+    inline static void SetItem(TDictObj *self, Any key, Any value) { (*self)[key] = value; }
+    inline static Any DelItem(TDictObj *self, Any key) { return Erase(self, key); }
     inline static Any GetKey(TDictObj *self, int64_t i) { return IterStateMut{self, i}.At().first; }
     inline static Any GetValue(TDictObj *self, int64_t i) { return IterStateMut{self, i}.At().second; }
     inline static int64_t Advance(TDictObj *self, int64_t i) { return IterStateMut{self, i}.Add().i; }
@@ -298,17 +300,19 @@ inline DictBase::KVPair *DictBase::Accessor<TDictObj>::TryInsertOrLookup(TDictOb
 }
 
 template <typename TDictObj> //
-inline void DictBase::Accessor<TDictObj>::Erase(TDictObj *self, const Any &key) {
+inline Any DictBase::Accessor<TDictObj>::Erase(TDictObj *self, const Any &key) {
   BlockIter iter = TSelf::Lookup(self, key);
   if (!iter.IsNone()) {
-    TSelf::Erase(self, iter.i);
-  } else {
-    MLC_THROW(KeyError) << key;
+    Any ret = static_cast<Any &>(iter.Data().second);
+    TSelf::_Erase(self, iter.i);
+    return ret;
   }
+  MLC_THROW(KeyError) << key;
+  MLC_UNREACHABLE();
 }
 
 template <typename TDictObj> //
-inline void DictBase::Accessor<TDictObj>::Erase(TDictObj *self, int64_t index) {
+inline void DictBase::Accessor<TDictObj>::_Erase(TDictObj *self, int64_t index) {
   DictBase *self_base = static_cast<DictBase *>(self);
   BlockIter iter = BlockIter::FromIndex(self_base, index);
   if (uint64_t offset = iter.Offset(); offset != 0) {
diff --git a/include/mlc/core/list.h b/include/mlc/core/list.h
@@ -157,7 +157,24 @@ struct UList : public ObjectRef {
       .FieldReadOnly("data", &MLCList::data)
       .StaticFn("__init__", FromAnyTuple)
       .MemFn("__str__", &UListObj::__str__)
-      .MemFn("__iter_at__", &::mlc::core::ListBase::Accessor<UListObj>::At);
+      .MemFn("__iter_at__", &::mlc::core::ListBase::Accessor<UListObj>::At)
+      .MemFn("_append", &UListObj::push_back)
+      .MemFn("_insert", [](UListObj *self, int64_t i, Any data) { self->insert(i, data); })
+      .MemFn("_extend",
+             [](int32_t num_args, const AnyView *args, Any *) {
+               if (!args[0].IsInstance<UListObj>()) {
+                 MLC_THROW(TypeError) << "First argument must be a list";
+               }
+               UListObj *self = args[0];
+               self->insert(self->size(), args + 1, args + num_args);
+             })
+      .MemFn("_pop",
+             [](UListObj *self, int64_t i) {
+               Any ret = self->operator[](i);
+               self->erase(i);
+               return ret;
+             })
+      .MemFn("_clear", &UListObj::clear);
 };
 
 template <typename T> struct ListObj : protected UListObj {
diff --git a/include/mlc/core/visitor.h b/include/mlc/core/visitor.h
@@ -15,6 +15,12 @@ namespace core {
 void ReportTypeFieldError(const char *type_key, MLCTypeField *field);
 
 template <typename Visitor> inline void VisitFields(Object *root, MLCTypeInfo *info, Visitor &&visitor) {
+  if (root == nullptr) {
+    MLC_THROW(ValueError) << "Root is nullptr";
+  }
+  if (info == nullptr) {
+    info = Lib::GetTypeInfo(root->GetTypeIndex());
+  }
   MLCTypeField *fields = info->fields;
   MLCTypeField *field = fields;
   for (; field->name != nullptr; ++field) {
@@ -90,6 +96,12 @@ template <typename Visitor> inline void VisitFields(Object *root, MLCTypeInfo *i
 }
 
 template <typename Visitor> inline void VisitStructure(Object *root, MLCTypeInfo *info, Visitor &&visitor) {
+  if (root == nullptr) {
+    MLC_THROW(ValueError) << "Root is nullptr";
+  }
+  if (info == nullptr) {
+    info = Lib::GetTypeInfo(root->GetTypeIndex());
+  }
   if (info->structure_kind == 0) {
     MLC_THROW(TypeError) << "Structure is not defined for type: " << info->type_key;
   }
diff --git a/include/mlc/printer/ast.h b/include/mlc/printer/ast.h
@@ -152,11 +152,11 @@ struct LiteralObj : public ::mlc::Object {
 }; // struct LiteralObj
 
 struct Literal : public ::mlc::printer::Expr {
-  static Literal Bool(bool value) { return Literal(mlc::List<ObjectPath>(), Any(value)); }
-  static Literal Int(int64_t value) { return Literal(mlc::List<ObjectPath>(), Any(value)); }
-  static Literal Str(mlc::Str value) { return Literal(mlc::List<ObjectPath>(), Any(value)); }
-  static Literal Float(double value) { return Literal(mlc::List<ObjectPath>(), Any(value)); }
-  static Literal Null() { return Literal(mlc::List<ObjectPath>(), Any()); }
+  static Literal Bool(bool value, List<ObjectPath> source_paths = {}) { return Literal(source_paths, Any(value)); }
+  static Literal Int(int64_t value, List<ObjectPath> source_paths = {}) { return Literal(source_paths, Any(value)); }
+  static Literal Str(mlc::Str value, List<ObjectPath> source_paths = {}) { return Literal(source_paths, Any(value)); }
+  static Literal Float(double value, List<ObjectPath> source_paths = {}) { return Literal(source_paths, Any(value)); }
+  static Literal Null(List<ObjectPath> source_paths = {}) { return Literal(source_paths, Any()); }
 
   MLC_DEF_OBJ_REF(MLC_EXPORTS, Literal, LiteralObj, ::mlc::printer::Expr)
       .Field("source_paths", &LiteralObj::source_paths)
diff --git a/include/mlc/printer/ir_printer.h b/include/mlc/printer/ir_printer.h
@@ -108,11 +108,26 @@ struct IRPrinterObj : public Object {
     return (*it).second->creator();
   }
 
-  Any operator()(const Optional<ObjectRef> &opt_obj, const ObjectPath &path) const {
-    if (!opt_obj.has_value()) {
-      return Literal::Null();
+  Any operator()(Any source, ObjectPath path) const {
+    if (source.type_index == kMLCNone) {
+      return Literal::Null({path});
     }
-    Node ret = ::mlc::Lib::IRPrint(opt_obj.value(), this, path);
+    if (source.type_index == kMLCBool) {
+      return Literal::Bool(source.operator bool(), {path});
+    }
+    if (source.type_index == kMLCInt) {
+      return Literal::Int(source.operator int64_t(), {path});
+    }
+    if (source.type_index == kMLCStr || source.type_index == kMLCRawStr) {
+      return Literal::Str(source.operator Str(), {path});
+    }
+    if (source.type_index == kMLCFloat) {
+      return Literal::Float(source.operator double(), {path});
+    }
+    if (source.type_index < kMLCStaticObjectBegin) {
+      MLC_THROW(ValueError) << "Unsupported type: " << source;
+    }
+    Node ret = ::mlc::Lib::IRPrint(source.operator Object *(), this, path);
     ret->source_paths->push_back(path);
     return ret;
   }
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,14 +1,14 @@
 [project]
 name = "mlc-python"
-version = "0.1.1"
+dynamic = ["version"]
 dependencies = [
     'numpy >= 1.22',
     'ml-dtypes >= 0.1',
     'Pygments>=2.4.0',
     'colorama',
     'setuptools ; platform_system == "Windows"',
 ]
-description = ""
+description = "Python-first Development for AI Compilers"
 requires-python = ">=3.9"
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -38,10 +38,14 @@ dev = [
 ]
 
 [build-system]
-requires = ["scikit-build-core>=0.9.8", "cython"]
+requires = ["scikit-build-core>=0.9.8", "cython", "setuptools-scm"]
 build-backend = "scikit_build_core.build"
 
+[tool.setuptools_scm]
+version_file = "python/mlc/_version.py"
+
 [tool.scikit-build]
+metadata.version.provider = "scikit_build_core.metadata.setuptools_scm"
 build.targets = ["mlc_py", "mlc_registry"]
 build.verbose = true
 cmake.build-type = "RelWithDebInfo"
diff --git a/python/mlc/__init__.py b/python/mlc/__init__.py
@@ -2,3 +2,9 @@
 from ._cython import Ptr, Str
 from .core import DataType, Device, Dict, Error, Func, List, Object, ObjectPath, json_loads, typing
 from .dataclasses import PyClass, c_class, py_class
+
+try:
+    from ._version import __version__, __version_tuple__  # type: ignore[import-not-found]
+except ImportError:
+    __version__ = version = "0.0.0.dev0"
+    __version_tuple__ = version_tuple = (0, 0, 0, "dev0", "ga1bb7a5c.d20230415")
diff --git a/python/mlc/core/dict.py b/python/mlc/core/dict.py
@@ -14,6 +14,12 @@
 T = TypeVar("T")
 
 
+class _UnspecifiedType: ...
+
+
+Unspecified = _UnspecifiedType()
+
+
 class DictMeta(MetaNoSlots, ABCMeta): ...
 
 
@@ -76,6 +82,28 @@ def items(self) -> ItemsView[K, V]:
     def __getitem__(self, key: K) -> V:
         return Dict._C(b"__getitem__", self, key)
 
+    def __setitem__(self, key: K, value: V) -> None:
+        Dict._C(b"__setitem__", self, key, value)
+
+    def __delitem__(self, key: K) -> None:
+        Dict._C(b"__delitem__", self, key)
+
+    def pop(self, key: K, default: T | _UnspecifiedType = Unspecified) -> V | T:
+        try:
+            return Dict._C(b"__delitem__", self, key)
+        except KeyError:
+            if default is Unspecified:
+                raise
+            assert not isinstance(default, _UnspecifiedType)
+            return default
+
+    def setdefault(self, key: K, default: V | None = None) -> V | None:
+        try:
+            return self[key]
+        except KeyError:
+            self[key] = default  # type: ignore[assignment]
+            return default
+
     # Additional methods required by the Mapping ABC
     def __contains__(self, key: object) -> bool:
         try:
diff --git a/python/mlc/core/list.py b/python/mlc/core/list.py
@@ -27,25 +27,46 @@ def __len__(self) -> int:
         return self.size
 
     @overload
-    def __getitem__(self, index: int) -> T: ...
+    def __getitem__(self, i: int) -> T: ...
 
     @overload
-    def __getitem__(self, index: slice) -> Sequence[T]: ...
-
-    def __getitem__(self, index: int | slice) -> T | Sequence[T]:
-        if isinstance(index, int):
-            length = len(self)
-            if not -length <= index < length:
-                raise IndexError(f"list index out of range: {index}")
-            if index < 0:
-                index += length
-            return List._C(b"__iter_at__", self, index)
-        elif isinstance(index, slice):
+    def __getitem__(self, i: slice) -> Sequence[T]: ...
+
+    def __getitem__(self, i: int | slice) -> T | Sequence[T]:
+        if isinstance(i, int):
+            i = _normalize_index(i, len(self))
+            return List._C(b"__iter_at__", self, i)
+        elif isinstance(i, slice):
             # Implement slicing
-            start, stop, step = index.indices(len(self))
+            start, stop, step = i.indices(len(self))
             return List([self[i] for i in range(start, stop, step)])
         else:
-            raise TypeError(f"list indices must be integers or slices, not {type(index).__name__}")
+            raise TypeError(f"list indices must be integers or slices, not {type(i).__name__}")
 
     def __iter__(self) -> Iterator[T]:
         return iter(self[i] for i in range(len(self)))
+
+    def insert(self, i: int, x: T) -> None:
+        i = _normalize_index(i, len(self) + 1)
+        return List._C(b"_insert", self, i, x)
+
+    def append(self, x: T) -> None:
+        return List._C(b"_append", self, x)
+
+    def pop(self, i: int = -1) -> T:
+        i = _normalize_index(i, len(self))
+        return List._C(b"_pop", self, i)
+
+    def clear(self) -> None:
+        return List._C(b"_clear", self)
+
+    def extend(self, iterable: Iterable[T]) -> None:
+        return List._C(b"_extend", self, *iterable)
+
+
+def _normalize_index(i: int, length: int) -> int:
+    if not -length <= i < length:
+        raise IndexError(f"list index out of range: {i}")
+    if i < 0:
+        i += length
+    return i
diff --git a/python/mlc/printer/__init__.py b/python/mlc/printer/__init__.py
@@ -2,4 +2,14 @@
 
 from . import ast
 from .ast import PrinterConfig
-from .ir_printer import DefaultFrame, IRPrinter, print_python, to_python
+from .ir_printer import (
+    Bool,
+    DefaultFrame,
+    Float,
+    Int,
+    IRPrinter,
+    None_,
+    Str,
+    print_python,
+    to_python,
+)
diff --git a/python/mlc/printer/ir_printer.py b/python/mlc/printer/ir_printer.py
diff --git a/tests/python/test_core_dict.py b/tests/python/test_core_dict.py
diff --git a/tests/python/test_core_list.py b/tests/python/test_core_list.py
diff --git a/tests/python/test_printer_ir_printer.py b/tests/python/test_printer_ir_printer.py
