Move patching to patch file

Author: staadecker

switch_model/solve.py

@@ -22,7 +22,6 @@
 )
 from switch_model.upgrade import do_inputs_need_upgrade, upgrade_inputs
 from switch_model.tools.graph.cli_graph import main as graph_main
-from switch_model.utilities.load_data import patch_to_allow_loading
 from switch_model.utilities.patches import patch_pyomo
 from switch_model.utilities.results_info import save_info, add_info, ResultsInfoSection
 
@@ -91,9 +90,6 @@ def debug(type, value, tb):
         # Patch pyomo if needed, to allow reconstruction of expressions.
         # This must be done before the model is constructed.
         patch_pyomo()
-        # This allows us to use input_file= when defining a Set or Param
-        patch_to_allow_loading(Set)
-        patch_to_allow_loading(Param)
 
         # Define the model
         model = create_model(modules, args=args)

switch_model/utilities/patches.py

@@ -5,6 +5,9 @@
 import pyomo.version
 from pyomo.core.base.misc import _robust_sort_keyfcn
 
+from switch_model.utilities.load_data import patch_to_allow_loading
+from pyomo.environ import Set, Param
+
 
 def fixed_robust_sort_keyfcn(self, val, use_key=True):
     """Generate a tuple ( str(type_name), val ) for sorting the value.
@@ -55,6 +58,10 @@ def patch_pyomo():
     from pyomo.core.base.misc import _robust_sort_keyfcn
     setattr(_robust_sort_keyfcn, "__call__", fixed_robust_sort_keyfcn)
 
+    # Patch Set and Param to allow specifying input file location (via input_file="...")
+    patch_to_allow_loading(Set)
+    patch_to_allow_loading(Param)
+
 def replace_method(class_ref, method_name, new_source_code):
     """
     Replace specified class method with a compiled version of new_source_code.