From 1baa906e519b246ce83bf804c77e3a7fb443c26d Mon Sep 17 00:00:00 2001
From: JCGoran <JCGoran@users.noreply.github.com>
Date: Tue, 27 Feb 2024 22:18:36 +0100
Subject: [PATCH] Fix Python lib issue with sympy solver (#1165)

* Fix for Python sympy solver

* Add corresponding test

* Update requirements

---------

Co-authored-by: Luc Grosheintz <luc.grosheintz@gmail.com>
---
 nmodl/__init__.py                | 34 ++++++++++++++++++++++++++++++++
 requirements.txt                 |  1 +
 setup.py                         |  3 ++-
 test/unit/pybind/test_visitor.py | 15 ++++++++++++++
 4 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/nmodl/__init__.py b/nmodl/__init__.py
index d0dfe8d4a..dbc53cbe9 100644
--- a/nmodl/__init__.py
+++ b/nmodl/__init__.py
@@ -1,7 +1,41 @@
+import os
+import sys
+
+if sys.version_info >= (3, 9):
+    from importlib.resources import files
+else:
+    from importlib_resources import files
+
+from find_libpython import find_libpython
+
+# try to add libpython*.so path to environment if not already set
+try:
+    os.environ["NMODL_PYLIB"] = os.environ.get(
+        "NMODL_PYLIB",
+        find_libpython(),
+    )
+except TypeError as exc:
+    raise RuntimeError(
+        "find_libpython was unable to find the Python library on this platform; "
+        "please make sure that the Python library is installed correctly\n"
+        "You can also try to manually set the NMODL_PYLIB environmental variable "
+        "to the Python library"
+    ) from exc
+
+# add nmodl home to environment (i.e. necessary for nrnunits.lib) if not
+# already set
+# `files` will automatically raise a `ModuleNotFoundError`
+os.environ["NMODLHOME"] = os.environ.get(
+    "NMODLHOME",
+    str(files("nmodl") / ".data"),
+)
+
+
 try:
     # Try importing but catch exception in case bindings are not available
     from ._nmodl import NmodlDriver, to_json, to_nmodl  # noqa
     from ._nmodl import __version__
+
     __all__ = ["NmodlDriver", "to_json", "to_nmodl"]
 except ImportError:
     print("[NMODL] [warning] :: Python bindings are not available")
diff --git a/requirements.txt b/requirements.txt
index 8816849e0..673379a57 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,3 +6,4 @@ sympy
 numpy
 find_libpython
 scikit-build
+importlib_resources;python_version<"3.9"
diff --git a/setup.py b/setup.py
index 5059f5150..095d6e656 100644
--- a/setup.py
+++ b/setup.py
@@ -76,7 +76,8 @@ def run(self, *args, **kwargs):
 install_requirements = [
     "PyYAML>=3.13",
     "sympy>=1.3",
-    "find_libpython"
+    "find_libpython",
+    "importlib_resources;python_version<'3.9'",
 ]
 
 
diff --git a/test/unit/pybind/test_visitor.py b/test/unit/pybind/test_visitor.py
index f402ea0c7..14329ba63 100644
--- a/test/unit/pybind/test_visitor.py
+++ b/test/unit/pybind/test_visitor.py
@@ -109,3 +109,18 @@ def visit_range_var(self, node):
 }
 """
     assert str(modast) == one_var_after
+
+
+def test_sympy_conductance_visitor():
+    """
+    Make sure NMODL sets the correct env variables to be able to run the sympy visitor
+    """
+    program = """NEURON {
+        USEION na READ ena WRITE ina
+        RANGE gna
+    }
+    BREAKPOINT {
+        ina = gna*(v - ena)
+    }"""
+    driver = nmodl.NmodlDriver()
+    visitor.SympyConductanceVisitor().visit_program(driver.parse_string(program))