在 Python 中调用 Mojo：基于 PythonModuleBuilder 构建高性能扩展模块完整指南

在 Python 中调用 Mojo基于 PythonModuleBuilder 构建高性能扩展模块完整指南【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo如果你已经拥有一个 Python 项目并且希望用 Mojo 的高性能计算能力加速其中关键的性能瓶颈并不需要把整个项目用 Mojo 重写。Mojo 提供了一套内建的 Python 扩展模块机制让你只需要把性能关键的部分用 Mojo 编写然后像普通 Python 模块一样import并调用。本篇指南以 Mojo 官方手册中「Calling Mojo from Python」章节Mojo/docs/site/manual/python/mojo-from-python.mdx为核心结合仓库内对应的可运行示例Mojo/docs/site/code/manual/python/mojo-from-python与 stdlib 的PythonModuleBuilder源码实现完整讲解从最小可运行示例、类型绑定、对象构造、方法暴露、关键字/可变参数到将 Python 代码移植为 Mojo 的实战策略帮助你快速掌握在 Python 中调用 Mojo 的全部技术要点。注意从 Python 调用 Mojo 目前属于Beta 实验特性处于早期开发阶段API 与使用体验预计会有较多变化文档也在持续完善中。文末列出了当前已知限制请在使用时留意。最小可运行示例从 Python 导入 Mojo 模块项目结构与最小代码考虑如下项目结构main.py是 Python 入口程序mojo_module.mojo是包含供 Python 调用的函数的 Mojo 源码project ├── main.py └── mojo_module.mojo假设我们想要一个接受 Python 值作为参数的 Mojo 函数比如计算阶乘。Mojo 一侧的初始写法为def factorial(py_obj: PythonObject) raises - PythonObject: var n Int(pypy_obj) return std.math.factorial(n)注意在使用PythonObject前需要导入相关模块。仓库中的完整示例文件Mojo/docs/site/code/manual/python/mojo-from-python/mojo_module.mojo实际是这样导入的import std.math from std.os import abort from std.python import PythonObject from std.python.bindings import PythonModuleBuilder在 Python 一侧我们期望这样调用import mojo_module print(mojo_module.factorial(5))为什么必须声明 PyInit_ 入口函数在从 Python 调用 Mojo 函数之前必须先让 Python 知道这个模块的存在。Python 加载mojo_module时会寻找名为PyInit_mojo_module()的函数如果文件叫foo.mojo则对应寻找PyInit_foo()。在PyInit_module()内部必须使用PythonModuleBuilder声明所有可从 Python 调用的 Mojo 函数与类型。因此完整的 Mojo 模块代码是from std.python import PythonObject from std.python.bindings import PythonModuleBuilder from std import math from std.os import abort export def PyInit_mojo_module() abi(C) - PythonObject: try: var m PythonModuleBuilder(mojo_module) m.def_functionfactorial return m.finalize() except e: abort(String(error creating Python Mojo module:, e)) def factorial(py_obj: PythonObject) raises - PythonObject: # 若 py_obj 无法转换为 Mojo Int这里会抛出异常 var n Int(pypy_obj) return math.factorial(n)在 Python 一侧需要把包含mojo_module.mojo的目录加入 Python 路径然后先导入mojo.importer再用普通的import语句加载 Mojo 代码import mojo.importer import mojo_module print(mojo_module.factorial(5))运行python main.py输出120仓库中对应的可运行示例正是这种结构main.py中先import mojo.importer再import mojo_module见 main.py并通过 Bazel 定义了modular_py_binary名为factorial与对应的modular_run_binary_testfactorial_test来执行这个例子见 BUILD.bazel。背后的工作原理Python 支持一种标准的「Python 扩展模块」机制使得编译型语言如 Mojo、C、C、Rust能够以直观的方式被 Python 调用。具体来说Python 扩展模块就是一个定义了合适的PyInit_*()函数的动态库。Mojo 内建了定义 Python 扩展模块的功能而真正“神奇”的部分发生在导入的mojo.importer模块中。导入 Mojo 代码后观察文件系统会发现多了一个__mojocache__目录内部有一个动态库.so文件project ├── main.py ├── mojo_module.mojo └── __mojocache__ └── mojo_module.hash-ABC123.so加载mojo.importer会注册 Python 的 Mojo import hook它在后台查找与导入模块名匹配的.mojo文件如果找到就使用mojo build --emit shared-lib将其编译为动态库产物存放在__mojocache__中并且只在缓存过期时通常是 Mojo 源文件发生变化才重新编译。提示清理缓存构建产物__mojocache__目录中应该只包含派生产物删除其中的内容是安全的。下次导入 Mojo 模块时需要的产物会自动重建。导出函数的 abi 约定export函数必须用显式的abi效果声明其调用约定。在 Python 扩展模块中唯一需要导出的是PyInit_module入口点并且它必须使用abi(C)export def PyInit_mojo_module() abi(C) - PythonObject: ...这是因为 CPython 运行时会直接跨 C 边界定位并调用PyInit_module所以它必须暴露 C 调用约定。abi(C)函数不能标记为raises这也是为什么上面的例子在函数体内捕获所有错误并通过abort处理而不是向上传播。相比之下你用模块构建器注册的函数、方法和初始化器def_function、def_method、def_py_init等完全不需要export你通过引用把它们传给构建器Mojo 会自动生成 CPython 真正调用的 C 包装器。这个包装器以 Mojo 调用约定调用你的函数并把任何抛出的错误转换为 Python 异常所以像factorial这样的注册函数可以放心地标记为raises。从 stdlib 源码bindings.mojo可以看到def_function要求函数签名满足参数类型均为PythonObject或末尾带var **kwargs: PythonObject返回类型必须是PythonObject或None不带关键字参数的函数通过 CPython 的METH_FASTCALL约定注册接受关键字参数的函数则使用METH_VARARGS | METH_KEYWORDS。Bindings 特性把 Mojo 类型与函数暴露给 Python绑定 Mojo 类型使用PythonModuleBuilder可以把任意 Mojo 类型绑定给 Python 使用。例如fieldwise_init struct Person(Movable, Writable): var name: String var age: Int export def PyInit_person_module() abi(C) - PythonObject: try: var mb PythonModuleBuilder(person_module) var person_type mb.add_typePerson except e: abort(error creating Mojo module)调用add_type()会返回一个PythonTypeBuilder随后可以用它绑定类型构造函数见下文「在 Python 中构造 Mojo 对象」和方法。任何通过PythonTypeBuilder绑定的 Mojo 类型其对应的 Pythontype对象都会被全局注册从而启用两个特性用PythonObject(allocPerson(..))构造包装 Mojo 值的 Python 对象供 Python 侧使用用python_obj.downcast_value_ptr[Person]()进行向下转换。注意要被绑定到 Python 使用Mojo 类型必须实现Writable。特定绑定特性还需要额外 trait自定义初始化器def_py_init要求Movable默认初始化器def_init_defaultable则同时要求Defaultable与Movable。仅绑定类型还不够还需要告诉 Python 如何与 Mojo 类型交互——先从在 Python 中构造 Mojo 对象实例开始。在 Python 中构造 Mojo 对象可以通过def_py_init()把 Mojo 初始化器声明为 Python 兼容的对象初始化器。例如export def PyInit_person_module() abi(C) - PythonObject: try: var mb PythonModuleBuilder(person_module) # highlight-start _ mb.add_typePerson.def_py_init[Person.py_init]() # highlight-end return mb.finalize() except e: abort(String(error creating Python Mojo module:, e)) fieldwise_init struct Person(Movable, Writable): var name: String var age: Int # highlight-start staticmethod def py_init( out self: Person, args: PythonObject, kwargs: PythonObject ) raises: # 校验参数个数 if len(args) ! 2: raise Error(Person() takes exactly 2 arguments) # 把 Python 参数转换为 Mojo 类型 var name String(args[0]) var age Int(args[1]) self Self(name, age) # highlight-end有了这个绑定就可以在 Python 中创建Person实例person person_module.Person(Sarah, 32) print(person)输出Person(nameSarah, age32)对于支持默认构造的类型可以使用更简单的def_init_defaultable()var counter_type m.add_typeCounter counter_type.def_init_defaultable[Counter]()这使 Python 代码可以无参数创建实例counter counter_module.Counter() # 创建 Counter()「构造器」与「初始化器」的区别在 Python 中对象构造横跨__new__()与__init__()两个方法__init__()严格来说是属性初始化器。但 Mojo struct 中没有__new__()方法所以我们始终把__init__()称为初始化器。将 Mojo 对象返回给 Python从 Python 调用的 Mojo 函数不仅需要能接受PythonObject参数还需要能返回新值有时甚至需要把 Mojo 原生值返回给 Python。这可以通过PythonObject(allocvalue)构造器实现def create_person() - PythonObject: var person Person(Sarah, 32) return PythonObject(allocperson^)警告如果所提供的 Mojo 对象类型之前没有通过PythonModuleBuilder.add_type()注册PythonObject(alloc...)会抛出异常。从 PythonObject 转换为 Mojo 值在处理PythonObject的 Mojo 代码中尤其是从 Python 调用的 Mojo 函数内通常期望参数是某个特定类型。把PythonObject变成 Mojo 原生值有两种方式转换Converting把 Python 对象转换为一个逻辑值相同的新构造的 Mojo 值由ConvertibleFromPythontrait 处理向下转换Downcasting把持有 Mojo 原生值的 Python 对象转换为指向该内部值的指针由PythonObject.downcast_value_ptr()处理。PythonObject 转换许多 Mojo 类型通过ConvertibleFromPythontrait 支持从等价 Python 类型直接转换# 给定一个 person克隆它并换一个名字 def create_person( name_obj: PythonObject, age_obj: PythonObject ) raises - PythonObject: # 转换失败时会抛出异常 var name String(name_obj) var age Int(age_obj) return PythonObject(allocPerson(name, age))从 Python 调用person mojo_module.create_person(John Smith)传入非法参数会导致运行时参数错误person mojo_module.create_person(42)从源码搜索可见ConvertibleFromPython目前已在若干基础类型上实现包括Int、Bool、SIMD与String等参见 int.mojo、bool.mojo、simd.mojo 与 string.mojo不过很多 stdlib 类型尚未实现该 trait。PythonObject 向下转换从PythonObject值向下转换到内部 Mojo 值def print_age(person_obj: PythonObject) raises: # 若 obj 不含 Mojo Person 类型的实例则抛出异常 var person person_obj.downcast_value_ptr[Person]() print(Person is, person[].age, years old)也支持通过向下转换进行不安全的修改。用户需要自行确保这个可变指针不与 Mojo 内部指向同一对象的其他指针发生别名def birthday(person_obj: PythonObject): var person person_obj.downcast_value_ptr[Person]() person[].age 1完全不做类型检查的向下转换可以用unchecked_downcast_value_ptrdef get_person(person_obj: PythonObject): var person person_obj.unchecked_downcast_value_ptr[Person]()不检查类型的向下转换可以消除类型检查的开销适合在已通过基准测试确认类型检查是瓶颈的紧内循环中使用。绑定方法绑定 Mojo 对象供 Python 使用时可以使用PythonTypeBuilder.def_method()把选定的方法暴露给 Python。目前暴露给 Python 的 Mojo 方法需要相对普通 Mojo 方法做一点修改必须是staticmethod且接收py_self: PythonObject或self_ptr: Pointer[Self]from std.python import PythonObject from std.python.bindings import PythonModuleBuilder from std.os import abort export def PyInit_mojo_module() abi(C) - PythonObject: try: var mb PythonModuleBuilder(mojo_module) # highlight-start _ mb.add_typePerson .def_methodPerson.get_name .def_methodPerson.set_age # highlight-end return mb.finalize() except e: abort(error creating Mojo module) struct Person(Writable): var name: String var age: Int # highlight-start staticmethod def get_name(py_self: PythonObject) raises - PythonObject: var self_ptr py_self.downcast_value_ptr[Self]() return self_ptr[].name staticmethod def set_age( self_ptr: Pointer[mutTrue, Self], new_age: PythonObject, ) raises: self_ptr[].age Int(new_age) # highlight-end def write_to(self, mut writer: Some[Writer]): tPerson({self.name}, {self.age}).write_to(writer)使用py_self: PythonObject可以访问 Mojo 对象实例所存放的完整PythonObject分配而一般情况下如果方法只是需要访问对象的字段使用self_ptr: Pointer[Self]可以减少样板代码。从 Python 调用的 Mojo 方法目前要求非标准的 self 类型这是受当前实现的限制未来版本的 Python Mojo bindings 会解除。绑定静态方法Python Mojo bindings 支持暴露 Python 风格的staticmethod通过PythonTypeBuilder.def_staticmethod()绑定。用def_staticmethod()声明的函数在 Python 中可以作为类型上的静态方法调用无需对象实例from std.python import PythonObject from std.python.bindings import PythonModuleBuilder from std.os import abort export def PyInit_mojo_module() abi(C) - PythonObject: try: var mb PythonModuleBuilder(mojo_module) # highlight-start mb.add_typePerson .def_staticmethodPerson.is_valid_age # highlight-end return mb.finalize() except e: abort(error creating Mojo module) struct Person(Writable): var name: String var age: Int # highlight-start staticmethod def is_valid_age(age_obj: PythonObject) raises - PythonObject: var age Int(age_obj) return 0 age 130 # highlight-end def write_to(self, mut writer: Some[Writer]): tPerson({self.name}, {self.age}).write_to(writer)在 Python 中调用绑定为静态方法的 Mojo 函数就是典型的 Python 静态方法调用from mojo_module import Person print(Person.is_valid_age(45)) # 输出 True print(Person.is_valid_age(-1)) # 输出 False关键字参数Mojo 的关键字参数有两种形式仅关键字参数def foo(*, x: Int)—— 目前 Python Mojo bindings不支持可变关键字参数def foo(var **kwargs: Int)—— 在 bindings 中支持但需要使用去糖形式def foo(kwargs: StringDict)**kwargs语法限制将在未来移除。可以用StringDict[PythonObject]作为最后一个参数来定义接受可变关键字参数的 Mojo 函数。简单示例import mojo_module result mojo_module.sum_kwargs_ints(a10, b20, c30) # 返回 60from std.collections import StringDict def sum_kwargs_ints(kwargs: StringDict[PythonObject]) raises - PythonObject: var total 0 for entry in kwargs.items(): total Int(entry.value) return PythonObject(total)关键字参数也支持跟在普通位置参数之后同时获取特定关键字参数就是对StringDict做字典查找from std.collections import StringDict def duration_in_seconds( hours_obj: PythonObject, minutes_obj: PythonObject, kwargs: StringDict[PythonObject] ) raises - PythonObject: var hours Int(hours_obj) var minutes Int(minutes_obj) var seconds Int(kwargs[seconds]) return hours * 3600 minutes * 60 seconds在这个例子中如果调用duration_in_seconds()时缺少必需的seconds命名参数会触发运行时异常from mojo_module import duration_in_seconds # 传入 hours 和 minutes缺少 seconds duration_in_seconds(4, 5) # ERROR: KeyError关键字参数在绑定顶层函数、方法和静态方法时均受支持。可变参数Python 与 Mojo 的可变参数通常写成def foo(*args: Int): ...但这个语法目前还不被 Python/Mojo bindings 支持因为用def_function()绑定的函数只支持固定元数fixed-arity。作为变通方案可以使用更底层的def_py_function()接口把接受可变数量参数的 Mojo 函数暴露给 Python参数个数校验由用户自己负责export def PyInit_mojo_module() abi(C) - PythonObject: try: var b PythonModuleBuilder(mojo_module) b.def_py_functioncount_args b.def_py_functionsum_args b.def_py_functionlookup def count_args(py_self: PythonObject, args_tuple: PythonObject) raises: return len(args_tuple) def sum_args(py_self: PythonObject, args_tuple: PythonObject) raises: var total args_tuple[0] for i in range(1, len(args_tuple)): total args_tuple[i] return total def lookup(py_self: PythonObject, args_tuple: PythonObject) raises: if len(args_tuple) ! 2 and len(args_tuple) ! 3: raise Error(lookup() expects 2 or 3 arguments) var collection args_tuple[0] var key args_tuple[1] try: return collection[key] except e: if len(args) 3: return args_tuple[2] else: raise e从 stdlib 源码可以看到def_py_function支持两类函数签名PyFunctionRaising与PyFunctionWithKeywordsRaising见 bindings.mojo最终都通过PyCFunction/PyCFunctionFast等 CPython 调用约定注册到模块中。将 Python 移植到 Mojo 的策略写出 Pythonic 的 Mojo在这种绑定思路下我们拥抱 Python 的灵活性不试图把PythonObject参数强制塞进 Mojo 强类型系统那个狭窄受限的空间而是直接写代码如果写错了就让它运行时抛异常。PythonObject的灵活性带来了一种独特的编程风格Python 代码几乎可以原样「移植」到 Mojo。def foo(x, y, z): x[y] int(z) x y z经验法则任何 Python 内建函数在 Mojo 中都应该可以通过Python.builtin()访问。def foo(x: PythonObject, y: PythonObject, z: PythonObject) - PythonObject: x[y] Python.int(z) x y z构建 Mojo 扩展模块的两种方式你可以通过以下方式创建并分发供 Python 使用的 Mojo 模块以源码文件形式分发通过 Python Mojo importer hook 按需编译。这种方式的优点是上手容易、项目结构简单并且编辑后导入的 Mojo 代码总是最新的。以预编译的 Python 扩展模块.so动态库分发使用命令编译mojo build mojo_module.mojo --emit shared-lib -o mojo_module.so这种方式的优点是可以手动指定其他任何必要的构建选项优化或调试标志、导入路径等为高级用户提供了绕过 Mojo import hook 抽象层的「逃生舱」。已知限制Python 与 Mojo 的互操作目标远大——Mojo 希望成为扩展 Python 的最佳方式——但该特性仍处于早期且活跃的开发阶段存在以下需要注意的限制未来会逐步解除关键字参数语法。目前从 Python 调用的 Mojo 函数只有在使用末尾的kwargs: StringDict[PythonObject]参数时才接受关键字参数原生**kwargs语法支持将在未来加入。Mojo 包依赖。依赖除 Mojo stdlib 之外其他包如 Modular Community 包渠道中的包的 Mojo 代码目前只在手动构建 Mojo 扩展模块时受支持因为 Mojo import hook 目前不支持为 Mojo 包依赖指定导入路径。属性Properties。计算属性的 getter 与 setter 目前不受支持。预期的类型转换。少数 Mojo 标准库类型通过实现ConvertibleFromPythontrait可以直接从等价的 Python 内建对象类型构造但许多 Mojo 标准库类型尚未实现该 trait如有需要可能得编写手动转换逻辑。参考资料Mojo 手册章节原文Calling Mojo from Python可运行示例目录Mojo/docs/site/code/manual/python/mojo-from-python含 mojo_module.mojo、main.py 与 BUILD.bazelstdlib 绑定实现bindings.mojoPythonModuleBuilder/PythonTypeBuilderPythonObject实现python_object.mojo转换 traitconversions.mojo【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考