加入 infinicore Python 包的 editable 安装支持

.gitignore

@@ -1,6 +1,7 @@
 # Xmake cache
 .xmake/
 build/
+python/infinicore/lib/*.so
 
 # MacOS Cache
 .DS_Store

python/infinicore/device.py

@@ -1,4 +1,4 @@
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class device:

python/infinicore/dtype.py

@@ -1,4 +1,4 @@
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class dtype:

python/infinicore/ops/matmul.py

@@ -1,7 +1,6 @@
+from infinicore.lib import _infinicore
 from infinicore.tensor import Tensor
 
-from .. import _infinicore
-
 
 def matmul(input, other, *, out=None):
     if out is None:

python/infinicore/ops/rearrange.py

@@ -1,7 +1,6 @@
+from infinicore.lib import _infinicore
 from infinicore.tensor import Tensor
 
-from .. import _infinicore
-
 
 def rearrange(input, other, *, out=None):
     if out is None:

python/infinicore/tensor.py

@@ -1,7 +1,6 @@
 import infinicore.device
 import infinicore.dtype
-
-from . import _infinicore
+from infinicore.lib import _infinicore
 
 
 class Tensor:

setup.py

@@ -1,8 +1,4 @@
-import glob
-import os
-import shutil
 import subprocess
-from pathlib import Path
 
 from setuptools import setup
 from setuptools.command.build import build
@@ -15,12 +11,5 @@ def run(self):
         subprocess.run(["xmake", "build", "-y", "_infinicore"])
         subprocess.run(["xmake", "install", "_infinicore"])
 
-        installation_dir = os.getenv("INFINI_ROOT", str(Path.home() / ".infini"))
-        lib_dir = os.path.join(installation_dir, "lib")
-        lib_path = glob.glob(os.path.join(lib_dir, "_infinicore.*"))[0]
-        package_dir = os.path.join(self.build_lib, "infinicore")
-        os.makedirs(package_dir, exist_ok=True)
-        shutil.move(lib_path, package_dir)
-
 
 setup(package_dir={"": "python"}, cmdclass={"build": Build})

xmake.lua

@@ -347,7 +347,31 @@ target("_infinicore")
     add_files("src/infinicore/op/*/*.cc")
     add_files("src/infinicore/pybind11/**.cc")
 
-    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+    set_installdir("python/infinicore")
+target_end()
+
+option("editable")
+    set_default(false)
+    set_showmenu(true)
+    set_description("Install the `infinicore` Python package in editable mode")
+option_end()
+
+target("infinicore")
+    set_kind("phony")
+
+    set_default(false)
+
+    add_deps("_infinicore")
+
+    on_install(function (target)
+        local pip_install_args = {}
+
+        if has_config("editable") then
+            table.insert(pip_install_args, "--editable")
+        end
+
+        os.execv("python", table.join({"-m", "pip", "install"}, pip_install_args, {"."}))
+    end)
 target_end()
 
 -- Tests