Update julia_import.py

pysr/julia_import.py

@@ -4,23 +4,12 @@
 from types import ModuleType
 from typing import cast
 
+from .julia_registry_helper import with_juliaregistrypref
 
-def get_registry_pref_env_var_name():
-    name = "JULIA_PKG_SERVER_REGISTRY_PREFERENCE"
-    return name
-
-
-def restore_julia_registry_pref(old_value):
-    name = get_registry_pref_env_var_name()
-    if old_value is None:
-        # Delete the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable that we set:
-        os.environ.pop(name)
-    else:
-        # Restore the original value of the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable:
-        os.environ[name] = old_value
+def import_juliacall():
+    import juliacall
     return None
 
-
 # Check if JuliaCall is already loaded, and if so, warn the user
 # about the relevant environment variables. If not loaded,
 # set up sensible defaults.
@@ -59,67 +48,15 @@ def restore_julia_registry_pref(old_value):
     # Deprecated; so just pass to juliacall
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
-if "JULIA_PKG_SERVER_REGISTRY_PREFERENCE" in os.environ:
-    old_julia_registry_pref = os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]
-    # If the user has set the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable, then
-    # we don't overwrite it.
-    pass
-else:
-    # If the user has not set the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable, then
-    # we set it temporarily, and unset it when we are done.
-    old_julia_registry_pref = None
-    os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
-
-try:
-    import juliacall
-except:
-    error_msg = """\
-        ERROR: Encountered a network error.
-        Are you behind a firewall, or are there network restrictions that would prevent access
-        to certain websites or domains?
-        Try setting the `{env_var_name}` environment variable to `conservative`.
-    """.format(
-        env_var_name=get_registry_pref_env_var_name()
-    )
-    if old_julia_registry_pref is not None:
-        print()
-        print(error_msg)
-        print()
-        restore_julia_registry_pref(old_julia_registry_pref)
-        # In this case, we know that the user had set JULIA_PKG_SERVER_REGISTRY_PREFERENCE,
-        # and thus we respect that value, and we don't override it.
-        # So, in this case, we just rethrow the caught exception:
-        raise  # rethrow the caught exception
-    else:
-        # In this case, the user had not set JULIA_PKG_SERVER_REGISTRY_PREFERENCE.
-        # So we had initially tried with JULIA_PKG_SERVER_REGISTRY_PREFERENCE=eager, but that
-        # resulted in an exception.
-        # So let us now try with JULIA_PKG_SERVER_REGISTRY_PREFERENCE=conservative
-        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "conservative"
-        try:
-            # Note: after changing the value of `JULIA_PKG_SERVER_REGISTRY_PREFERENCE`,
-            # you need to run `Pkg.Registry.update()` again. Otherwise the change will not take effect.
-            # `juliacall` will automatically do this for us, so we don't need to do it ourselves.
-            #
-            # See line 334 here:
-            # https://github.com/JuliaPy/pyjuliapkg/blob/3a2c66019f098c1ebf84f933a46e7ca70e82792b/src/juliapkg/deps.py#L334-L334
-            import juliacall
-        except:
-            print()
-            print(error_msg)
-            print()
-            # Now, we just rethrow the caught exception:
-            restore_julia_registry_pref(old_julia_registry_pref)
-            raise  # rethrow the caught exception
+# Run `import juliacall`, but inside our `with_juliaregistrypref()` wrapper:
+with_juliaregistrypref(import_juliacall)
 
 
 import juliacall
 from juliacall import AnyValue  # type: ignore
 from juliacall import VectorValue  # type: ignore
 from juliacall import Main as jl  # type: ignore
 
-restore_julia_registry_pref(old_julia_registry_pref)
-
 
 jl = cast(ModuleType, jl)