Repository: PlasmaControl/DESC

Issue:
Breaking changes to tests in Python 3.13
Python 3.13 breaks how we use the `inspect` module. After some debugging, it appears in Python 3.13 only the source code of objects that are functions as classified by `inspect.isfunction` can be retrieved. In particular, JITed functions can no longer have their source code inspected.

One test that fails is ``tests/test_data_index.py``.
To resolve that failure, one solution is, in that file, to replace
```for _, fun in inspect.getmembers(module, callable):```
with 
```for _, fun in inspect.getmembers(module, inspect.isfunction):```

and also unjit all compute functions in `desc/_compute`. It would be nice to fix this without unjitting the compute functions though.


Otherwise, the error message that appears is
```python
    def findsource(object):
        """Return the entire source file and starting line number for an object.
    
        The argument may be a module, class, method, function, traceback, frame,
        or code object.  The source code is returned as a list of all the lines
        in the file and the line number indexes a line in that list.  An OSError
        is raised if the source code cannot be retrieved."""
    
        file = getsourcefile(object)
        if file:
            # Invalidate cache if needed.
            linecache.checkcache(file)
        else:
            file = getfile(object)
            # Allow filenames in form of "<something>" to pass through.
            # `doctest` monkeypatches `linecache` module to enable
            # inspection, so let `linecache.getlines` to be called.
            if (not (file.startswith('<') and file.endswith('>'))) or file.endswith('.fwork'):
                raise OSError('source code not available')
    
        module = getmodule(object, file)
        if module:
            lines = linecache.getlines(file, module.__dict__)
        else:
            lines = linecache.getlines(file)
        if not lines:
            raise OSError('could not get source code')
    
        if ismodule(object):
            return lines, 0
    
        if isclass(object):
            try:
                lnum = vars(object)['__firstlineno__'] - 1
            except (TypeError, KeyError):
>               raise OSError('source code not available')
E               OSError: source code not available


```

Discussion:

Patch:
diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py
index 2beb08dad..70d712b6f 100644
--- a/desc/objectives/_generic.py
+++ b/desc/objectives/_generic.py
@@ -1,6 +1,5 @@
 """Generic objectives that don't belong anywhere else."""
 
-import inspect
 import re
 
 import numpy as np
@@ -11,7 +10,7 @@ from desc.compute.utils import _compute as compute_fun
 from desc.compute.utils import _parse_parameterization, get_profiles, get_transforms
 from desc.grid import QuadratureGrid
 from desc.optimizable import OptimizableCollection
-from desc.utils import errorif, jaxify, parse_argname_change, setdefault
+from desc.utils import errorif, getsource, jaxify, parse_argname_change, setdefault
 
 from .linear_objectives import _FixedObjective
 from .objective_funs import _Objective, collect_docs
@@ -557,7 +556,7 @@ class ObjectiveFromUser(_Objective):
 
         def get_vars(fun):
             pattern = r"data\[(.*?)\]"
-            src = inspect.getsource(fun)
+            src = getsource(fun)
             variables = re.findall(pattern, src)
             variables = list({s.strip().strip("'").strip('"') for s in variables})
             return variables
diff --git a/desc/utils.py b/desc/utils.py
index f6a2c3f8c..0069b5c09 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -1,6 +1,7 @@
 """Utility functions, independent of the rest of DESC."""
 
 import functools
+import inspect
 import operator
 import warnings
 from itertools import combinations_with_replacement, permutations
@@ -965,3 +966,10 @@ def ensure_tuple(x):
     if isinstance(x, list):
         return tuple(x)
     return (x,)
+
+
+def getsource(obj):
+    """Get source code for an object, allowing for partials etc."""
+    if isinstance(obj, functools.partial):
+        return getsource(obj.func)
+    return inspect.getsource(obj)


Test patch:
diff --git a/tests/test_axis_limits.py b/tests/test_axis_limits.py
index bc3bbe144..e85cec9ff 100644
--- a/tests/test_axis_limits.py
+++ b/tests/test_axis_limits.py
@@ -6,7 +6,6 @@ If the limit has yet to be derived, add it to the ``not_implemented_limits`` set
 """
 
 import functools
-import inspect
 
 import numpy as np
 import pytest
@@ -18,7 +17,7 @@ from desc.examples import get
 from desc.grid import LinearGrid
 from desc.integrals import surface_integrals_map
 from desc.objectives import GenericObjective, ObjectiveFunction
-from desc.utils import dot, errorif
+from desc.utils import dot, errorif, getsource
 
 # Unless mentioned in the source code of the compute function, the assumptions
 # made to compute the magnetic axis limit can be reduced to assuming that these
@@ -366,7 +365,7 @@ def _reverse_mode_unsafe_names():
 
     def get_source(name):
         return "".join(
-            inspect.getsource(
+            getsource(
                 data_index["desc.equilibrium.equilibrium.Equilibrium"][name]["fun"]
             ).split("def ")[1:]
         )
diff --git a/tests/test_data_index.py b/tests/test_data_index.py
index ebb3df247..51d9cce55 100644
--- a/tests/test_data_index.py
+++ b/tests/test_data_index.py
@@ -8,12 +8,12 @@ import pytest
 import desc.compute
 from desc.compute import data_index
 from desc.compute.data_index import _class_inheritance
-from desc.utils import errorif
+from desc.utils import errorif, getsource
 
 
 def _get_matches(fun, pattern, ignore_comments=True):
     """Return all matches of ``pattern`` in source code of function ``fun``."""
-    src = inspect.getsource(fun)
+    src = getsource(fun)
     if ignore_comments:
         # remove any decorator functions
         src = src.partition("def ")[2]
@@ -27,7 +27,7 @@ def _get_matches(fun, pattern, ignore_comments=True):
 def _get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"):
     """Get parameterization of thing computed by function ``fun``."""
     pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")')
-    decorator = inspect.getsource(fun).partition("def ")[0]
+    decorator = getsource(fun).partition("def ")[0]
     matches = pattern.findall(decorator)
     # if list was found, split strings in list, else string was found so get that
     matches = [match[0].split(",") if match[0] else [match[1]] for match in matches]
@@ -75,8 +75,9 @@ def test_data_index_deps():
     for module_name, module in inspect.getmembers(desc.compute, inspect.ismodule):
         if module_name[0] == "_":
             # JITed functions are not functions according to inspect,
-            # so just check if callable.
-            for _, fun in inspect.getmembers(module, callable):
+            # so just check if callable and in the right module
+            filt = lambda x: callable(x) and x.__module__ == module.__name__
+            for _, fun in inspect.getmembers(module, filt):
                 register_name = _get_matches(fun, pattern_name, ignore_comments=False)
                 if not register_name:
                     continue