From 84585e4da751bc5942cdaf25c2414c5d5d850e27 Mon Sep 17 00:00:00 2001
From: Brinton Eldridge <brinton1516@gmail.com>
Date: Tue, 3 Sep 2024 08:38:59 -0500
Subject: [PATCH] reverted _eval_step_func, rm get_name_args, edited
 get_num_args

---
 python/simulation.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/python/simulation.py b/python/simulation.py
index 2ed0c475c..013cc0ccb 100644
--- a/python/simulation.py
+++ b/python/simulation.py
@@ -3,6 +3,7 @@
 """
 
 import functools
+import inspect
 import math
 import numbers
 import os
@@ -96,18 +97,12 @@ def get_num_args(func):
     return (
         2
         if isinstance(func, Harminv) or isinstance(func, PadeDFT)
+        else func.__code__.co_argcount - 1
+        if inspect.ismethod(func)
         else func.__code__.co_argcount
     )
 
 
-def get_name_args(func):
-    return (
-        ("sim", "todo")
-        if isinstance(func, Harminv) or isinstance(func, PadeDFT)
-        else func.__code__.co_varnames[: func.__code__.co_argcount]
-    )
-
-
 def vec(*args):
     try:
         # Check for vec(x, [y, [z]])
@@ -5004,16 +4999,14 @@ def _combine(sim, todo):
 
 def _eval_step_func(sim, func, todo):
     num_args = get_num_args(func)
-    name_args = get_name_args(func)
-    self_count = int("self" in name_args)
 
-    if num_args not in {1 + self_count, 2 + self_count}:
-        raise ValueError(f"Step function '{func.__name__}' requires 1 or 2 arguments")
-    elif num_args == 2 + self_count:
-        func(sim, todo)
-    elif num_args == 1 + self_count:
+    if num_args != 1 and num_args != 2:
+        raise ValueError(f"Step function '{func.__name__}'' requires 1 or 2 arguments")
+    elif num_args == 1:
         if todo == "step":
             func(sim)
+    elif num_args == 2:
+        func(sim, todo)
 
 
 def _when_true_funcs(cond, *step_funcs):