From f03c8f9307c19ac90ec2767ee6d46f1d5fcce384 Mon Sep 17 00:00:00 2001
From: nihui <nihuini@tencent.com>
Date: Sun, 4 Feb 2024 21:33:40 +0800
Subject: [PATCH] build universal pnnx wheel and pnnx entrypoint script (#5331)

---
 .ci/pnnx.yml                       |  6 ++--
 .gitignore                         |  3 +-
 tools/pnnx/python/README.md        |  6 ++--
 tools/pnnx/python/pnnx/__init__.py |  5 ++++
 tools/pnnx/{ => python}/setup.py   |  7 +++--
 tools/pnnx/python/setup.py.i       | 45 ------------------------------
 6 files changed, 16 insertions(+), 56 deletions(-)
 rename tools/pnnx/{ => python}/setup.py (96%)
 delete mode 100644 tools/pnnx/python/setup.py.i

diff --git a/.ci/pnnx.yml b/.ci/pnnx.yml
index 16602627c43..082dc8fac2f 100644
--- a/.ci/pnnx.yml
+++ b/.ci/pnnx.yml
@@ -132,7 +132,7 @@ jobs:
       run: |
         export PYTHONUSERBASE=${{ci.workspace}}/torch-${{matrix.torch-version}}
         export PNNX_WHEEL_WITHOUT_BUILD=ON
-        cd tools/pnnx
-        cp build/src/pnnx python/pnnx/
+        cd tools/pnnx/python
+        cp ../build/src/pnnx pnnx/
         python3 setup.py install --user
-        pytest python/tests/
+        pytest tests
diff --git a/.gitignore b/.gitignore
index 2c71aee0332..cd69c526f19 100644
--- a/.gitignore
+++ b/.gitignore
@@ -54,10 +54,9 @@ __pycache__
 *.pyd
 *.egg-info/
 python/setup.py
-tools/pnnx/python/setup.py
 
 # Clangd
 .cache/
 
 # Xmake
-.xmake/
\ No newline at end of file
+.xmake/
diff --git a/tools/pnnx/python/README.md b/tools/pnnx/python/README.md
index 964d8fa53ff..4aabf9c699f 100644
--- a/tools/pnnx/python/README.md
+++ b/tools/pnnx/python/README.md
@@ -42,7 +42,7 @@ conda install pytorch
 ```
 3. install
 ```bash
-cd /pathto/ncnntools/pnnx
+cd /pathto/ncnntools/pnnx/python
 python setup.py install
 ```
 
@@ -61,8 +61,8 @@ python setup.py install
 
 ## Tests
 ```bash
-cd /pathto/ncnn/tools/pnnx
-pytest python/tests/
+cd /pathto/ncnn/tools/pnnx/python
+pytest tests
 ```
 
 ## Usage
diff --git a/tools/pnnx/python/pnnx/__init__.py b/tools/pnnx/python/pnnx/__init__.py
index 0ba69277d7a..e33f5c55abe 100644
--- a/tools/pnnx/python/pnnx/__init__.py
+++ b/tools/pnnx/python/pnnx/__init__.py
@@ -13,7 +13,10 @@
 # specific language governing permissions and limitations under the License.
 
 import os
+import sys
 import platform
+import subprocess
+
 EXEC_DIR_PATH = os.path.dirname(os.path.abspath(__file__))
 if platform.system() == 'Linux' or platform.system() == "Darwin":
     EXEC_PATH = EXEC_DIR_PATH + "/pnnx"
@@ -31,3 +34,5 @@
 except:
     pass
 
+def pnnx():
+    raise SystemExit(subprocess.call([EXEC_PATH] + sys.argv[1:], close_fds=False))
diff --git a/tools/pnnx/setup.py b/tools/pnnx/python/setup.py
similarity index 96%
rename from tools/pnnx/setup.py
rename to tools/pnnx/python/setup.py
index dcbb99922b3..bb3ca546ed1 100644
--- a/tools/pnnx/setup.py
+++ b/tools/pnnx/python/setup.py
@@ -156,10 +156,11 @@ def build_extension(self, ext):
     ],
     license="BSD-3",
     python_requires=">=3.7",
-    packages=find_packages("python"),
+    packages=find_packages(),
     package_data={"pnnx": ["pnnx", "pnnx.exe"]},
-    package_dir={"": "python"},
+    package_dir={"": "."},
     install_requires=requirements,
-    ext_modules=[CMakeExtension("pnnx")],
+    ext_modules=None if PNNX_WHEEL_WITHOUT_BUILD == 'ON' else [CMakeExtension("pnnx", "..")],
     cmdclass={"build_ext": CMakeBuild},
+    entry_points={"console_scripts": ["pnnx=pnnx:pnnx"]},
 )
diff --git a/tools/pnnx/python/setup.py.i b/tools/pnnx/python/setup.py.i
deleted file mode 100644
index 5c3c72e876d..00000000000
--- a/tools/pnnx/python/setup.py.i
+++ /dev/null
@@ -1,45 +0,0 @@
-import sys
-from setuptools import setup, find_packages
-
-try:
-    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
-
-    class bdist_wheel(_bdist_wheel):
-        def finalize_options(self):
-            _bdist_wheel.finalize_options(self)
-            self.root_is_pure = False
-
-except ImportError:
-    bdist_wheel = None
-
-if sys.version_info < (3, 0):
-    sys.exit("Sorry, Python < 3.0 is not supported")
-
-requirements = ["torch"]
-
-setup(
-        name="pnnx",
-        version="${PACKAGE_VERSION}",
-        author="nihui",
-        author_email="nihuini@tencent.com",
-        description="pnnx is an open standard for PyTorch model interoperability.",
-        url="https://github.com/Tencent/ncnn/tree/master/tools/pnnx",
-        classifiers=[
-            "Programming Language :: C++",
-            "Programming Language :: Python :: 3",
-            "Programming Language :: Python :: 3.7",
-            "Programming Language :: Python :: 3.8",
-            "Programming Language :: Python :: 3.9",
-            "Programming Language :: Python :: 3.10",
-            "Programming Language :: Python :: 3.11",
-            "License :: OSI Approved :: BSD License",
-            "Operating System :: OS Independent",
-        ],
-        license="BSD-3",
-        python_requires=">=3.6",
-        packages=find_packages(),
-        package_dir={"": "."},
-        package_data={"pnnx": ["pnnx${PYTHON_MODULE_PREFIX}${PYTHON_MODULE_EXTENSION}"]},
-        install_requires=requirements,
-        cmdclass={"bdist_wheel": bdist_wheel},
-)
\ No newline at end of file