From 4b82cb22f7a7eeeac524d91da7364abca98c620c Mon Sep 17 00:00:00 2001
From: Mehdi Ataei <mehdi.ataei@autodesk.com>
Date: Thu, 16 Jan 2025 16:07:04 -0500
Subject: [PATCH] Added packaging

---
 .gitignore      |  8 ++++++
 py/neon/gate.py | 21 ++++++++++++--
 py/setup.py     | 75 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 3 deletions(-)
 create mode 100644 py/setup.py

diff --git a/.gitignore b/.gitignore
index 1d06ae5a..5acf9ffe 100644
--- a/.gitignore
+++ b/.gitignore
@@ -65,3 +65,11 @@ py/neon/tool/Neon.log
 py/neon/Neon.log
 
 py/tests/Neon.log
+
+py/neon/liblibNeonPy.so
+
+py/dist/
+
+py/build/
+
+py/neon.egg-info/
diff --git a/py/neon/gate.py b/py/neon/gate.py
index 83b122c8..156db920 100644
--- a/py/neon/gate.py
+++ b/py/neon/gate.py
@@ -1,7 +1,7 @@
 import ctypes
 import os
 import warp as wp
-
+import platform
 
 class Gate(object):
     def __init__(self):
@@ -9,8 +9,23 @@ def __init__(self):
         # get the path of this python file
         current_file_path = os.path.abspath(__file__)
         # get the directory containing the script
-        lib_path = os.path.dirname(current_file_path) + "/../../cmake-build-debug/libNeonPy/liblibNeonPy.so"
-        # move up two folders with respec to script_dir
+        # Determine platform and architecture
+        current_platform = platform.system().lower()
+        machine = platform.machine().lower()
+        
+        if current_platform.startswith("linux"):
+            platform_key = f"linux_{machine}"
+            lib_extension = ".so"
+        elif current_platform.startswith("darwin"):
+            platform_key = f"macos_{machine}"
+            lib_extension = ".dylib" 
+        elif current_platform.startswith("windows"):
+            platform_key = f"windows_{machine}"
+            lib_extension = ".dll"
+        else:
+            raise RuntimeError(f"Unsupported platform: {current_platform}")
+            
+        lib_path = os.path.join(os.path.dirname(current_file_path), "lib", platform_key, f"liblibNeonPy{lib_extension}")
 
         try:
             self.lib =    ctypes.CDLL(lib_path)
diff --git a/py/setup.py b/py/setup.py
new file mode 100644
index 00000000..00b184f2
--- /dev/null
+++ b/py/setup.py
@@ -0,0 +1,75 @@
+import os
+import shutil
+import sys
+import platform
+from setuptools import setup, find_packages
+from setuptools.dist import Distribution
+from setuptools.command.build_py import build_py as build_py_orig
+
+class BinaryDistribution(Distribution):
+    """Distribution that includes binary components."""
+    def has_ext_modules(self):
+        return True
+
+class CustomBuildPy(build_py_orig):
+    """Custom build command to copy the appropriate shared library."""
+    def run(self):
+        # Determine platform and architecture
+        current_platform = sys.platform
+        machine = platform.machine().lower()
+
+        # Map sys.platform to your binaries directory
+        if current_platform.startswith("linux"):
+            platform_key = "linux_" + machine
+            lib_extension = ".so"
+        elif current_platform.startswith("darwin"):
+            platform_key = "macos_" + machine
+            lib_extension = ".dylib"
+        elif current_platform.startswith("win"):
+            platform_key = "windows_" + machine
+            lib_extension = ".dll"
+        else:
+            raise RuntimeError(f"Unsupported platform: {current_platform}")
+
+        # Define the path to the shared library. TODO: This is a hack to get the path to the shared library. We need to find a better way to do this for multiple platforms.
+        source_so = os.path.abspath(os.path.join(os.path.dirname(__file__), f'../cmake-build-debug/libNeonPy/liblibNeonPy{lib_extension}'))
+
+        # Define the destination directory within the build directory
+        destination_dir = os.path.join(self.build_lib, 'neon', 'lib', platform_key)
+        destination_so = os.path.join(destination_dir, os.path.basename(source_so))
+
+        # Ensure the destination directory exists
+        os.makedirs(destination_dir, exist_ok=True)
+
+        # Copy the shared library to the destination directory
+        try:
+            shutil.copy2(source_so, destination_so)
+            print(f"Copied {source_so} to {destination_so}")
+        except IOError as e:
+            print(f"Error copying {source_so} to {destination_so}: {e}")
+            raise e
+
+        # Continue with the standard build process
+        super().run()
+
+setup(
+    name="neon",
+    version="0.1.0",
+    packages=find_packages(),
+    package_data={
+        "neon": [
+            "**/*.h",             # Include all .h files recursively
+        ]
+    },
+    include_package_data=True,
+    distclass=BinaryDistribution,
+    cmdclass={
+        'build_py': CustomBuildPy,
+    },
+    python_requires=">=3.10", # This should be the minimum version of python that warp supports.
+    # TODO: We need to add warp to the requirements later. Currently we're using a custom build of warp.
+    install_requires=[
+        "numpy>=2.0",
+        "nvtx"
+    ],
+)