From f7ab7e41cdc14eb60262ebd82818360916a15314 Mon Sep 17 00:00:00 2001
From: cchall <chall@radiasoft.net>
Date: Mon, 1 Aug 2022 13:02:57 -0600
Subject: [PATCH] Move sirepo import to only be local

---
 rsopt/configuration/setup.py         | 11 ++++++-----
 rsopt/package_data/shifter_sirepo.py |  1 +
 rsopt/pkcli/pack.py                  |  5 +++--
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/rsopt/configuration/setup.py b/rsopt/configuration/setup.py
index c8b2b6f..4a9b314 100644
--- a/rsopt/configuration/setup.py
+++ b/rsopt/configuration/setup.py
@@ -4,7 +4,6 @@
 import pickle
 import typing
 import pathlib
-import sirepo.lib
 from rsopt import util
 from rsopt.codes import _TEMPLATED_CODES
 from rsopt import _SETUP_SCHEMA
@@ -68,7 +67,7 @@ def _validate_execution_type(key):
         return False
 
 
-def _shifter_parse_model(name: str, input_file: str, ignored_files: list) -> sirepo.lib.SimData or None:
+def _shifter_parse_model(name: str, input_file: str, ignored_files: list) -> typing.Type['sirepo.lib.SimData'] or None:
     # Sidesteps the difficulty of Sirepo install on NERSC by running a script that parses to the Sirepo model
     import shlex
     from subprocess import Popen, PIPE
@@ -149,7 +148,8 @@ def templated(cls):
 
     @classmethod
     def parse_input_file(cls, input_file: str, shifter: str,
-                         ignored_files: typing.Optional[typing.List[str]] = None) -> sirepo.lib.SimData or None:
+                         ignored_files: typing.Optional[typing.List[str]] = None) -> typing.Type['sirepo.lib.SimData'] \
+                                                                                     or None:
 
         if shifter:
             # Must pass a list to ignored_files here since it is sent to subprocess
@@ -295,7 +295,8 @@ def generate_input_file(self, kwarg_dict, directory):
         for k in dict_item_str.keys():
             kwarg_dict.pop(k)
 
-        output_template = template.render(dict_item=kwarg_dict, dict_item_str=dict_item_str, full_input_file_path=self.setup['input_file'],
+        output_template = template.render(dict_item=kwarg_dict, dict_item_str=dict_item_str,
+                                          full_input_file_path=self.setup['input_file'],
                                           function=self.setup['function'])
 
         file_path = os.path.join(directory, _PARALLEL_PYTHON_RUN_FILE)
@@ -527,4 +528,4 @@ def generate_input_file(self, kwarg_dict, directory):
     'opal': Opal,
     'user': User,
     'genesis': Genesis
-}
\ No newline at end of file
+}
diff --git a/rsopt/package_data/shifter_sirepo.py b/rsopt/package_data/shifter_sirepo.py
index 8b3733a..1f24cf8 100644
--- a/rsopt/package_data/shifter_sirepo.py
+++ b/rsopt/package_data/shifter_sirepo.py
@@ -1,3 +1,4 @@
+"""Script is called inside Shifter container with Sirepo installed and returns parsed model information"""
 import pickle
 import sys
 import sirepo.lib
diff --git a/rsopt/pkcli/pack.py b/rsopt/pkcli/pack.py
index 72d6638..1f4b8e0 100644
--- a/rsopt/pkcli/pack.py
+++ b/rsopt/pkcli/pack.py
@@ -1,10 +1,8 @@
 import rsopt.parse as parse
-import sirepo.sim_data  # Can't run import sirepo and call sirepo.sim_data. ??
 import modulefinder
 import tarfile
 import pathlib
 import os
-from sirepo.template import lattice
 
 
 def _get_local_modules(script_name):
@@ -32,6 +30,9 @@ def _get_processing(job):
 
 
 def _get_file_list_from_model(model, code_name, input_file_path):
+    import sirepo.sim_data
+    from sirepo.template import lattice
+
     # Use same path as configuration file when packing
     model_path = pathlib.Path(input_file_path).parents[0]
     # Get input and lattice if exists