From 2538d674d0b1c506d9ad8e02c371865cc512e5fb Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Wed, 4 Dec 2024 19:52:59 -0500
Subject: [PATCH] By default, use the `eager` flavor of the Julia registry from
 the Julia Pkg Servers (but try falling back to `conservative` if an exception
 is encountered)

---
 pysr/julia_import.py | 69 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 4d7b9150..41607ac6 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -4,6 +4,23 @@
 from types import ModuleType
 from typing import cast
 
+
+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
+    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.
@@ -42,10 +59,62 @@
     # 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:
+            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
+
+
+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)