如果您现有的 Python 项目能够从 Mojo 的高性能计算中受益,那么您无需用 Mojo 重写整个项目。相反,您可以只在 Mojo 中编写性能关键部分的代码,然后从 Python 中调用它。
在Python
为了说明从 Python 调用 Mojo 是什么样子,我们将从一个简单的示例开始,然后深入了解它的工作原理以及目前的可能性。
考虑具有以下结构的项目:
project├── 🐍 main.py
└── 🔥 mojo_module.mojo
主入口点是一个名为的 Python 程序main.py
,Mojo 代码包含从 Python 调用的函数。
例如,假设我们希望 Mojo 函数将 Python 值作为参数:
🔥 mojo_module.mojo
fn factorial(py_obj: PythonObject) raises -> Python var n = Int(py_obj) return math.factorial(n)
我们希望从 Python 中像这样调用它:
🐍 main.py
import mojo_moduleprint(mojo_module.factorial(5))
但是,在我们从 Python 调用 Mojo 函数之前,我们必须声明它,以便 Python 知道它的存在。
因为 Python 正在尝试加载mojo_module
,所以它会查找名为 的函数 PyInit_mojo_module()
。(如果我们的文件名为foo.mojo
,则该函数为PyInit_foo()
。)在 中PyInit_mojo_module()
,我们必须使用 声明所有可从 Python 调用的 Mojo 函数和类型 PythonModuleBuilder
。
因此完整的 Mojo 代码如下所示:
🔥 mojo_module.mojo
from python import PythonObjectfrom python.bindings import PythonModuleBuilderimport mathfrom os import abort@exportfn PyInit_mojo_module() -> PythonObject: try: var m = PythonModuleBuilder("mojo_module") m.def_function[factorial]("factorial", docstring="Compute n!") return m.finalize() except e: return abort[PythonObject](String("error creating Python Mojo module:", e))fn factorial(py_obj: PythonObject) raises -> PythonObject: # Raises an exception if `py_obj` is not convertible to a Mojo `Int`. var n = Int(py_obj) return math.factorial(n)
from python import PythonObjectfrom python.bindings import PythonModuleBuilderimport mathfrom os import abort@exportfn PyInit_mojo_module() -> PythonObject: try: var m = PythonModuleBuilder("mojo_module") m.def_function[factorial]("factorial", docstring="Compute n!") return m.finalize() except e: return abort[PythonObject](String("error creating Python Mojo module:", e))fn factorial(py_obj: PythonObject) raises -> PythonObject: # Raises an exception if `py_obj` is not convertible to a Mojo `Int`. var n = Int(py_obj) return math.factorial(n)
在 Python 方面,我们目前需要一些样板代码来使其工作(但这很快就会改善):
🐍 main.py
import max._mojo.mojo_importerimport osimport syssys.path.insert(0, "")os.environ["MOJO_PYTHON_LIBRARY"] = ""import mojo_moduleprint(mojo_module.factorial(5))
就这样!试试吧:
python main.py
120
工作原理
Python 支持一种称为Python 扩展模块的标准机制,该机制使编译型语言(例如 Mojo、C、C++ 或 Rust)能够以直观的方式从 Python 调用自身。具体来说,Python 扩展模块只是一个定义了合适函数的动态库PyInit_*()
。
Mojo 内置了定义 Python 扩展模块的功能。特殊功能在max._mojo.mojo_importer
我们导入的模块中实现。
如果我们在 Python 导入 Mojo 代码后查看文件系统,我们会注意到有一个新__mojocache__
目录,里面有动态库(.so
)文件:
project├── main.py├── mojo_module.mojo└── __mojocache__ └── mojo_module.hash-ABC123.so
加载时max._mojo.mojo_importer
会加载我们的 Python Mojo导入钩子,它会在后台查找与导入模块名称匹配的.mojo
(或) 文件,如果找到,则使用 编译它以生成静态库。生成的文件存储在 中,并且仅在其过期时(通常是 Mojo 源文件更改时)重建。.🔥
mojo build --emit shared-lib
__mojocache__
现在我们已经了解了如何在 Python 中使用 Mojo 的基础知识,让我们深入研究可用的功能以及如何利用它们通过 Mojo 加速您的 Python。
绑定 Mojo类型
所有 Mojo 类型都可以绑定到 Python 中使用。要将 Mojo 类型暴露给 Python,它必须实现 TypeIdentifiable
相应的特性。对于一个简单的 Mojo 类型,它可能如下所示:
🔥 Mojo
struct Person(TypeIdentifiable, ...): var name: String var age: Int # Unique name under which the type object is stored. # Eventually this will be a compiler provided unique type ID. alias TYPE_ID = "mojo_module.Person"
目前类型还必须实现Movable
、Defaultable
和 Representable
才能绑定到 Python 以供使用。
这使得可以使用以下方式绑定类型PythonModuleBuilder.add_type[Person]()
:
🔥 Mojo
var mb = PythonModuleBuilder("mojo_module")mb.add_type[Person]("Person")
任何使用 a 绑定的 Mojo 类型PythonTypeBuilder
都将使生成的 Python“类型”对象被全局注册,从而实现两个功能:
在Mojo
Python Mojo 绑定目前不支持__init__()
接受参数的 Mojo 方法。不过,可以使用构造新对象的自由函数(如下所示)来解决这个问题。
从 Python 调用的 Mojo 函数不仅需要能够接受 PythonObject
值作为参数,还需要能够返回新值。有时,它们甚至需要能够将 Mojo 原生值返回给 Python。这可以通过使用PythonObject(alloc=<value>)
构造函数来实现。
示例如下:
🔥 Mojo
fn create_person() -> PythonObject: var person = Person("Sarah", 32) return PythonObject(alloc=person^)
警告
PythonObject(alloc=...)
如果提供的 Mojo 对象类型之前未使用 注册,则会引发异常 PythonModuleBuilder.add_type()
。
PythonObject
Mojo价值观
在任何处理 的 Mojo 代码中 PythonObject
,尤其是在从 Python 调用的 Mojo 函数中,通常需要特定类型的参数。特定类型和愿望
有两种情况PythonObject
可以将 a “转换”为原生 Mojo 值:
PythonObject
转化
绑定初始化器。当前的一个限制是,非默认的 Mojo 类型__init__()
方法无法绑定到 Python 调用。除了展示参数转换之外,此示例还展示了如何使用顶级函数构造 Mojo 类型的实例并将其返回给 Python。
许多 Mojo 类型支持通过以下特征直接从等效的 Python 类型进行转换ConvertibleFromPython
:
🔥 Mojo
fn create_person( name_obj: PythonObject, age_obj: PythonObject) raises -> PythonObject: # These conversions will raise an exception if they fail var name = String(name_obj) var age = Int(age_obj) return PythonObject(alloc=Person(name, age))
可以使用以下方式从 Python 调用:
🐍 Python
person = mojo_module.create_person("John Smith", 42)
传递无效参数将导致类型错误:
🐍 Python
# TODO: What is the exact error message this emits today?person = mojo_module.create_person([1, 2, 3], {"foo": 4})
PythonObject
沮丧
从PythonObject
值向下转换为内部 Mojo 值:
🔥 Mojo
fn print_age(person_obj: PythonObject): # Raises if `obj` does not contain an instance of the Mojo `Person` type. var person = person_obj.downcast_value_ptr[Person]() # TODO(MSTDL-1581): # var person = Pointer[Person](downcast_value=person_obj) print("Person is", person[].age, "years old")
还支持通过向下转型实现不安全的可变。用户需要确保此可变指针不会在 Mojo 中成为指向同一对象的任何其他指针的别名:
🔥 Mojo
fn birthday(person_obj: PythonObject): var person = person_obj.downcast_value_ptr[Person]() # TODO: # var person = Pointer[Person](unsafe_unique_downcast=person_obj) person[].age += 1
完全不受检查的向下转型(不进行类型检查)可以使用以下方法完成:
🔥 Mojo
fn get_person(person_obj: PythonObject): var person = person_obj.unchecked_downcast_value_ptr[Person]() # TODO: # var person = Pointer[Person](unchecked_downcast_value=person_obj)
在使用 Mojo 优化紧密的内部循环时,可以使用未经检查的向下转换来消除开销,并且您已经进行基准测试和测量,类型检查向下转换是一个重大瓶颈。
在这种绑定方法中,我们利用了 Python 的灵活性,避免尝试将PythonObject
参数转换为 Mojo 类型系统的严格限制的强类型空间,而是只编写一些代码并让它在运行时引发异常(如果出现错误)。
灵活性PythonObject
实现了独特的编程风格,其中 Python 代码只需相对较少的更改即可“移植”到 Mojo。
🐍 Python
def foo(x, y, z):
x[y] = int(z)
x = y + z
经验法则:任何 Python 内置函数都应该可以在 Mojo 中使用 访问 Python.<builtin>()
。
🔥 Mojo
fn foo(x: PythonObject, y: PythonObject, z: PythonObject) -> PythonObject: x[y] = Python.int(z) x = y + z x.attr = z
Python Mojo 绑定目前本身不支持关键字参数,但可以使用一种简单的模式将它们提供给库的用户,即使用 Python 包装函数通过字典将关键字参数传递到 Mojo。
此模式的一个简单示例如下:
🐍 Python
import mojo_moduledef supports_kwargs(pos, *, kw1 = None, kw2 = None): mojo_module.supports_kwargs(pos, { "kw1": kw1, "kw2": kw2})
🔥 Mojo
fn supports_kwargs(pos: PythonObject, kwargs: PythonObject) raises: var kw1 = kwargs["kw1"] var kw2 = kwargs["kw2"]
由于关键字参数的验证和默认值在 Python 包装函数中处理,调用者将获得预期的标准参数错误。Mojo 代码保持简洁,因为获取关键字参数只需进行简单的字典查找。
使用 绑定函数时 PythonModuleBuilder.def_function()
,仅支持固定参数数量的函数。要将接受可变数量参数的 Mojo 函数公开给 Python,可以使用低级 def_py_function()
接口,这将验证所提供参数的数量。
🔥 Mojo
@exportfn PyInit_mojo_module() -> PythonObject: try: var b = PythonModuleBuilder("mojo_module") b.def_py_function[count_args]("count_args") b.def_py_function[sum_args]("sum_args") b.def_py_function[lookup]("lookup")fn count_args(py_self: PythonObject, args: TypedPythonObject["Tuple"]): return len(args)fn sum_args(py_self: PythonObject, args: TypedPythonObject["Tuple"]): var total = args[0] for i in range(1, len(args)): total += args[i] return totalfn lookup(py_self: PythonObject, args: TypedPythonObject["Tuple"]) raises: if len(args) != 2 and len(args) != 3: raise Error("lookup() expects 2 or 3 arguments") var collection = args[0] var key = args[1] try: return collection[key] except e: if len(args) == 3: return args[2] else: raise e
构建 Mojo 扩展模块
您可以通过以下方式为 Python 创建和分发 Mojo 模块:
作为源文件,使用 Python Mojo 导入器钩子按需编译。
这种方法的优点是易于上手,并使您的项目结构保持简单,同时确保您导入的 Mojo 代码在您进行编辑后始终保持最新。
作为预构建的 Python 扩展模块.so
动态库,使用以下方式编译:
$ mojo build mojo_module.mojo --emit shared-lib -o mojo_module.so
这样做的好处是您可以手动指定任何其他必要的构建选项(优化或调试标志、导入路径等),为高级用户提供从 Mojo 导入钩子抽象中“退出”的出口。