Fix Python interop (modular#517)

zhongcai1 · Aug 25, 2023 · a9dfe70 · a9dfe70
1 parent 5d65c38
commit a9dfe70
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 11 deletions.
diff --git a/examples/hello_interop.mojo b/examples/hello_interop.mojo
@@ -26,14 +26,11 @@ def main():
     print("Hello Mojo 🔥!")
     for x in range(9, 0, -3):
         print(x)
-    alias test_dir = "/root/mojo-examples"
     try:
-        Python.add_to_path(test_dir)
-        let test_module: PythonObject = Python.import_module("simple_interop")
-        if test_module:
-            _ = test_module.test_interop_func()
-        else:
-            print("Could not locate module: ", test_module)
+        Python.add_to_path(".")
+        Python.add_to_path("./examples")
+        let test_module = Python.import_module("simple_interop")
+        test_module.test_interop_func()
     except e:
         print(e.value)
-        pass
+        print("could not find module simple_interop")
diff --git a/examples/matmul.mojo b/examples/matmul.mojo
@@ -66,16 +66,16 @@ struct Matrix:
 fn run_matmul_python(M: Int, N: Int, K: Int) -> Float64:
     var gflops: Float64 = 0.0
     let python = Python()
-    alias test_dir = "/root/mojo-examples"
     try:
-        Python.add_to_path(test_dir)
+        Python.add_to_path(".")
+        Python.add_to_path("./examples")
         let pymatmul_module: PythonObject = Python.import_module("pymatmul")
         if pymatmul_module:
             gflops = pymatmul_module.benchmark_matmul_python(
                 M, N, K
             ).to_float64()
         else:
-            print("Python matmul module not found")
+            print("pymatmul module not found")
     except e:
         print(e.value)
         pass

diff --git a/examples/pymatmul.py b/examples/pymatmul.py
@@ -13,6 +13,14 @@
 
 # Simple program demonstrating a naive matrix multiplication in Python
 
+import importlib
+import sys
+import subprocess
+
+if not importlib.find_loader("numpy"):
+    print("Numpy not found, installing...")
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy"])
+
 import numpy as np
 from timeit import timeit
 

diff --git a/examples/simple_interop.py b/examples/simple_interop.py
@@ -14,6 +14,14 @@
 
 # Simple python program to test interop
 
+import importlib
+import sys
+import subprocess
+
+if not importlib.find_loader("numpy"):
+    print("Numpy not found, installing...")
+    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy"])
+
 import numpy as np
 from timeit import timeit