Merge pull request #348 from xpsi-group/347-python-3-version-needs-fr…

…om-custom-to-find-the-custom-file Module generator bug fixed to correctly import Custom files
xpsi-group · Dec 8, 2023 · b34efb5 · b34efb5
2 parents b86ea37 + 341cf60
commit b34efb5
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 18 deletions.
diff --git a/examples/examples_module_generator/_auto_modules/main.py b/examples/examples_module_generator/_auto_modules/main.py
@@ -855,16 +855,30 @@ def str_to_bool(x):
 from xpsi import HotRegions
 
 print('Rank reporting: %d' % xpsi._rank)
-from CustomInstrument import CustomInstrument
-from CustomSignal import CustomSignal
-from CustomInterstellar import CustomInterstellar
 
-try:
-    from CustomPhotosphere import CustomPhotosphere
-except ImportError:
-    from xpsi import Photosphere as CustomPhotosphere
+if __name__ == '__main__':
+    from CustomInstrument import CustomInstrument
+    from CustomSignal import CustomSignal
+    from CustomInterstellar import CustomInterstellar
+
+    try:
+        from CustomPhotosphere import CustomPhotosphere
+    except ImportError:
+        from xpsi import Photosphere as CustomPhotosphere
+
+    from CustomPrior import CustomPrior
+
+else:
+    from .CustomInstrument import CustomInstrument
+    from .CustomSignal import CustomSignal
+    from .CustomInterstellar import CustomInterstellar
+
+    try:
+        from .CustomPhotosphere import CustomPhotosphere
+    except ImportError:
+        from xpsi import Photosphere as CustomPhotosphere
 
-from CustomPrior import CustomPrior
+    from .CustomPrior import CustomPrior
 
 if args.main_import_statements is not None:
     for import_statement in args.main_import_statements:

diff --git a/xpsi/module_generator.py b/xpsi/module_generator.py
@@ -1405,16 +1405,29 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 module += (
 '''
-from {0} import CustomInstrument
-from {1} import CustomSignal
-from {2} import CustomInterstellar
+if __name__ == '__main__':
+    from {0} import CustomInstrument
+    from {1} import CustomSignal
+    from {2} import CustomInterstellar
 
-try:
-    from {3} import CustomPhotosphere
-except ImportError:
-    from xpsi import Photosphere as CustomPhotosphere
+    try:
+        from {3} import CustomPhotosphere
+    except ImportError:
+        from xpsi import Photosphere as CustomPhotosphere
+
+    from {4} import CustomPrior
+
+else:
+    from .{0} import CustomInstrument
+    from .{1} import CustomSignal
+    from .{2} import CustomInterstellar
+
+    try:
+        from .{3} import CustomPhotosphere
+    except ImportError:
+        from xpsi import Photosphere as CustomPhotosphere
 
-from {4} import CustomPrior
+    from .{4} import CustomPrior
 '''.format(args.custom_instrument_module,
            args.custom_signal_module,
            args.custom_interstellar_module,
@@ -1433,14 +1446,22 @@ def __call__(self, parser, namespace, values, option_string=None):
 if args.background_shared_class:
     module += (
     '''
-from {0} import CustomBackground
+if __name__ == '__main__':
+    from {0} import CustomBackground
+else:
+    from .{0} import CustomBackground
     '''.format(args.custom_background_module)
     )
 elif args.background_shared_class:
     for _instrument in args.instruments:
         module += (
         '''
-from {0} import {1}_CustomBackground
+if __name__ == '__main__':
+    from {0} import {1}_CustomBackground
+        '''.format(args.custom_background_module,
+                   _instrument)
+else:
+    from .{0} import {1}_CustomBackground
         '''.format(args.custom_background_module,
                    _instrument)
         )