From a40e8c77af9a890db944d5ef8545dd1cfc6b8460 Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Wed, 4 Dec 2024 19:52:59 -0500
Subject: [PATCH 01/14] 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_extensions.py      |   6 ++
 pysr/julia_import.py          |  14 +++++
 pysr/julia_registry_helper.py | 100 ++++++++++++++++++++++++++++++++++
 3 files changed, 120 insertions(+)
 create mode 100644 pysr/julia_registry_helper.py

diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index 72273f3e..9d00aeac 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -3,6 +3,7 @@
 from typing import Literal
 
 from .julia_import import Pkg, jl
+from .julia_registry_helper import with_juliaregistrypref
 from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
 
 
@@ -46,6 +47,11 @@ def isinstalled(uuid_s: str):
 
 
 def load_package(package_name: str, uuid_s: str) -> None:
+    with_juliaregistrypref(_load_package, package_name, uuid_s)
+    return None
+
+
+def _load_package(package_name: str, uuid_s: str) -> None:
     if not isinstalled(uuid_s):
         Pkg.add(name=package_name, uuid=uuid_s)
         Pkg.resolve()
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 4d7b9150..8187d2f0 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -4,6 +4,15 @@
 from types import ModuleType
 from typing import cast
 
+from .julia_registry_helper import with_juliaregistrypref
+
+
+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.
@@ -42,6 +51,11 @@
     # Deprecated; so just pass to juliacall
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
+# 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
diff --git a/pysr/julia_registry_helper.py b/pysr/julia_registry_helper.py
new file mode 100644
index 00000000..087308fa
--- /dev/null
+++ b/pysr/julia_registry_helper.py
@@ -0,0 +1,100 @@
+import os
+
+
+def get_juliaregistrypref_envvarname():
+    # We have this function so that we can define this string in one place, instead of needing
+    # to repeat ourselves in multiple places.
+    name = "JULIA_PKG_SERVER_REGISTRY_PREFERENCE"
+    return name
+
+
+def backup_juliaregistrypref():
+    name = get_juliaregistrypref_envvarname()
+    if name in os.environ:
+        old_value = 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_value = None
+
+        # We will set the JULIA_PKG_SERVER_REGISTRY_PREFERENCE to `eager`.
+        # Note: in upstream Julia/Pkg, the default value is `conservative`.
+        # Therefore, by setting it to `eager`, we are overriding the default.
+        # Because we are deviating from the default behavior that would normally be expected in Julia,
+        # I think that it's a good idea to print an informational message, so that the user is aware
+        # of the change being made.
+        info_msg = """\
+            INFO: Attempting to use the `eager` registry flavor of the Julia General registry from the Julia Pkg server.
+            If any errors are encountered, try setting the `{name}` environment variable to `conservative`.
+        """.format(
+            name=name  # TODO: get rid of this line?
+        )
+        print(info_msg)
+        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
+    return old_value
+
+
+def restore_juliaregistrypref(old_value):
+    name = get_juliaregistrypref_envvarname()
+    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
+
+
+def with_juliaregistrypref(f, *args):
+    name = get_juliaregistrypref_envvarname()
+    old_value = backup_juliaregistrypref()
+    try:
+        f(*args)
+    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 `{name}` environment variable to `conservative`.
+        """.format(
+            name=name  # TODO: get rid of this line?
+        )
+        if old_value is not None:
+            print()
+            print(error_msg)
+            print()
+
+            # 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:
+            restore_juliaregistrypref(
+                old_value
+            )  # Restore the old value BEFORE we re-throw.
+            raise  # Now, re-throw 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[name] = "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
+                f(args)
+            except:
+                print()
+                print(error_msg)
+                print()
+                # Now, we just rethrow the caught exception:
+                restore_juliaregistrypref(
+                    old_value
+                )  # Restore the old value BEFORE we re-throw.
+                raise  # Now, re-throw the caught exception.
+    return None

From 99b89c73af1dabf005b48b74cf9d62ffd222d6d6 Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Wed, 4 Dec 2024 21:35:58 -0500
Subject: [PATCH 02/14] Print our log messages to stderr (not stdout)

---
 pysr/julia_registry_helper.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/pysr/julia_registry_helper.py b/pysr/julia_registry_helper.py
index 087308fa..1d6b7eeb 100644
--- a/pysr/julia_registry_helper.py
+++ b/pysr/julia_registry_helper.py
@@ -1,4 +1,5 @@
 import os
+import sys
 
 
 def get_juliaregistrypref_envvarname():
@@ -32,7 +33,7 @@ def backup_juliaregistrypref():
         """.format(
             name=name  # TODO: get rid of this line?
         )
-        print(info_msg)
+        print(info_msg, file=sys.stderr)
         os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
     return old_value
 
@@ -63,9 +64,9 @@ def with_juliaregistrypref(f, *args):
             name=name  # TODO: get rid of this line?
         )
         if old_value is not None:
-            print()
-            print(error_msg)
-            print()
+            print("", file=sys.stderr)
+            print(error_msg, file=sys.stderr)
+            print("", file=sys.stderr)
 
             # 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.
@@ -89,9 +90,9 @@ def with_juliaregistrypref(f, *args):
                 # https://github.com/JuliaPy/pyjuliapkg/blob/3a2c66019f098c1ebf84f933a46e7ca70e82792b/src/juliapkg/deps.py#L334-L334
                 f(args)
             except:
-                print()
-                print(error_msg)
-                print()
+                print("", file=sys.stderr)
+                print(error_msg, file=sys.stderr)
+                print("", file=sys.stderr)
                 # Now, we just rethrow the caught exception:
                 restore_juliaregistrypref(
                     old_value

From cd571aa5e8946b1fd1814dd37c0e98e70f86bdc1 Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Fri, 6 Dec 2024 00:23:58 -0500
Subject: [PATCH 03/14] Try to fix the typecheck

---
 pysr/julia_import.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 8187d2f0..f61648fd 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -8,7 +8,7 @@
 
 
 def import_juliacall():
-    import juliacall
+    import juliacall # type: ignore
 
     return None
 

From 1e69b44ac8009de5ccd9f4740b2702c652dfec9d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 6 Dec 2024 05:24:07 +0000
Subject: [PATCH 04/14] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 pysr/julia_import.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index f61648fd..1bc21fdd 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -8,7 +8,7 @@
 
 
 def import_juliacall():
-    import juliacall # type: ignore
+    import juliacall  # type: ignore
 
     return None
 

From 9332eccf378e17a5c6ec4868b9b00c45279a3695 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 12:43:23 +0000
Subject: [PATCH 05/14] refactor: clean up registry helper code

---
 pysr/julia_registry_helper.py | 127 +++++++++++++---------------------
 1 file changed, 48 insertions(+), 79 deletions(-)

diff --git a/pysr/julia_registry_helper.py b/pysr/julia_registry_helper.py
index 1d6b7eeb..0b5bb7e2 100644
--- a/pysr/julia_registry_helper.py
+++ b/pysr/julia_registry_helper.py
@@ -1,101 +1,70 @@
-import os
-import sys
-
+"""Utilities for managing Julia registry preferences during package operations."""
 
-def get_juliaregistrypref_envvarname():
-    # We have this function so that we can define this string in one place, instead of needing
-    # to repeat ourselves in multiple places.
-    name = "JULIA_PKG_SERVER_REGISTRY_PREFERENCE"
-    return name
+import os
+import warnings
+from collections.abc import Callable
 
 
 def backup_juliaregistrypref():
-    name = get_juliaregistrypref_envvarname()
-    if name in os.environ:
-        old_value = 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_value = None
+    """Backup and potentially modify Julia registry preference.
 
-        # We will set the JULIA_PKG_SERVER_REGISTRY_PREFERENCE to `eager`.
-        # Note: in upstream Julia/Pkg, the default value is `conservative`.
-        # Therefore, by setting it to `eager`, we are overriding the default.
-        # Because we are deviating from the default behavior that would normally be expected in Julia,
-        # I think that it's a good idea to print an informational message, so that the user is aware
-        # of the change being made.
-        info_msg = """\
-            INFO: Attempting to use the `eager` registry flavor of the Julia General registry from the Julia Pkg server.
-            If any errors are encountered, try setting the `{name}` environment variable to `conservative`.
-        """.format(
-            name=name  # TODO: get rid of this line?
+    Sets JULIA_PKG_SERVER_REGISTRY_PREFERENCE to 'eager' if not already set.
+    Returns the original value for later restoration.
+    """
+    old_value = os.environ.get("JULIA_PKG_SERVER_REGISTRY_PREFERENCE", None)
+    if old_value is None:
+        warnings.warn(
+            "Attempting to use the `eager` registry flavor of the Julia "
+            "General registry from the Julia Pkg server.\n"
+            "    If any errors are encountered, try setting the "
+            "`JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment variable to `conservative`."
         )
-        print(info_msg, file=sys.stderr)
         os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
     return old_value
 
 
-def restore_juliaregistrypref(old_value):
-    name = get_juliaregistrypref_envvarname()
+def restore_juliaregistrypref(old_value: str | None):
+    """Restore the original Julia registry preference value."""
     if old_value is None:
-        # Delete the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable that we set:
-        os.environ.pop(name)
+        del os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]
     else:
-        # Restore the original value of the JULIA_PKG_SERVER_REGISTRY_PREFERENCE environment variable:
-        os.environ[name] = old_value
-    return None
+        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value
+
 
+def with_juliaregistrypref(f: Callable[..., None], *args):
+    """Execute function with modified Julia registry preference.
 
-def with_juliaregistrypref(f, *args):
-    name = get_juliaregistrypref_envvarname()
+    Temporarily modifies the registry preference to 'eager', falling back to
+    'conservative' if network errors occur. Restores original preference after
+    execution.
+
+    Parameters
+    ----------
+    f : Callable[..., None]
+        Function to execute. Should not return anything of importance.
+    *args : Any
+        Arguments to pass to the function.
+    """
     old_value = backup_juliaregistrypref()
     try:
         f(*args)
-    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 `{name}` environment variable to `conservative`.
-        """.format(
-            name=name  # TODO: get rid of this line?
+    except Exception as e:
+        error_msg = (
+            "ERROR: Encountered a network error.\n"
+            "    Are you behind a firewall, or are there network restrictions that would "
+            "prevent access to certain websites or domains?\n"
+            "    Try setting the `JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment "
+            "variable to `conservative`."
         )
         if old_value is not None:
-            print("", file=sys.stderr)
-            print(error_msg, file=sys.stderr)
-            print("", file=sys.stderr)
-
-            # 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:
-            restore_juliaregistrypref(
-                old_value
-            )  # Restore the old value BEFORE we re-throw.
-            raise  # Now, re-throw the caught exception.
+            warnings.warn(error_msg)
+            restore_juliaregistrypref(old_value)
+            raise e
         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[name] = "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
                 f(args)
-            except:
-                print("", file=sys.stderr)
-                print(error_msg, file=sys.stderr)
-                print("", file=sys.stderr)
-                # Now, we just rethrow the caught exception:
-                restore_juliaregistrypref(
-                    old_value
-                )  # Restore the old value BEFORE we re-throw.
-                raise  # Now, re-throw the caught exception.
-    return None
+            except Exception as e:
+                warnings.warn(error_msg)
+                restore_juliaregistrypref(old_value)
+                raise e

From 41cccad74cdba81a049b44dc22eef3b872e6e5d4 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 12:44:55 +0000
Subject: [PATCH 06/14] refactor: canonicalize filename

---
 pysr/julia_extensions.py                                     | 2 +-
 pysr/julia_import.py                                         | 2 +-
 pysr/{julia_registry_helper.py => julia_registry_helpers.py} | 0
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename pysr/{julia_registry_helper.py => julia_registry_helpers.py} (100%)

diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index 9d00aeac..cc926534 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -3,7 +3,7 @@
 from typing import Literal
 
 from .julia_import import Pkg, jl
-from .julia_registry_helper import with_juliaregistrypref
+from .julia_registry_helpers import with_juliaregistrypref
 from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
 
 
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 1bc21fdd..966100cd 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -4,7 +4,7 @@
 from types import ModuleType
 from typing import cast
 
-from .julia_registry_helper import with_juliaregistrypref
+from .julia_registry_helpers import with_juliaregistrypref
 
 
 def import_juliacall():
diff --git a/pysr/julia_registry_helper.py b/pysr/julia_registry_helpers.py
similarity index 100%
rename from pysr/julia_registry_helper.py
rename to pysr/julia_registry_helpers.py

From 4bfee9d0bacab27971a82367ff079521eccdece9 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:00:27 +0000
Subject: [PATCH 07/14] feat: try conservative registry first; only then switch

---
 pysr/julia_extensions.py       | 11 ++++----
 pysr/julia_import.py           |  4 ---
 pysr/julia_registry_helpers.py | 49 +++++++++++++++++++---------------
 3 files changed, 32 insertions(+), 32 deletions(-)

diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index cc926534..b8fad157 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -47,14 +47,13 @@ def isinstalled(uuid_s: str):
 
 
 def load_package(package_name: str, uuid_s: str) -> None:
-    with_juliaregistrypref(_load_package, package_name, uuid_s)
-    return None
+    if not isinstalled(uuid_s):
 
+        def _add_package():
+            Pkg.add(name=package_name, uuid=uuid_s)
+            Pkg.resolve()
 
-def _load_package(package_name: str, uuid_s: str) -> None:
-    if not isinstalled(uuid_s):
-        Pkg.add(name=package_name, uuid=uuid_s)
-        Pkg.resolve()
+        with_juliaregistrypref(_add_package)
 
     # TODO: Protect against loading the same symbol from two packages,
     #       maybe with a @gensym here.
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 966100cd..a1e5996e 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -10,8 +10,6 @@
 def import_juliacall():
     import juliacall  # type: ignore
 
-    return None
-
 
 # Check if JuliaCall is already loaded, and if so, warn the user
 # about the relevant environment variables. If not loaded,
@@ -51,11 +49,9 @@ def import_juliacall():
     # Deprecated; so just pass to juliacall
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
-# 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
diff --git a/pysr/julia_registry_helpers.py b/pysr/julia_registry_helpers.py
index 0b5bb7e2..fed82900 100644
--- a/pysr/julia_registry_helpers.py
+++ b/pysr/julia_registry_helpers.py
@@ -34,8 +34,8 @@ def restore_juliaregistrypref(old_value: str | None):
 def with_juliaregistrypref(f: Callable[..., None], *args):
     """Execute function with modified Julia registry preference.
 
-    Temporarily modifies the registry preference to 'eager', falling back to
-    'conservative' if network errors occur. Restores original preference after
+    First tries with existing registry preference. If that fails with a Julia registry error,
+    temporarily modifies the registry preference to 'eager'. Restores original preference after
     execution.
 
     Parameters
@@ -45,26 +45,31 @@ def with_juliaregistrypref(f: Callable[..., None], *args):
     *args : Any
         Arguments to pass to the function.
     """
-    old_value = backup_juliaregistrypref()
     try:
         f(*args)
-    except Exception as e:
-        error_msg = (
-            "ERROR: Encountered a network error.\n"
-            "    Are you behind a firewall, or are there network restrictions that would "
-            "prevent access to certain websites or domains?\n"
-            "    Try setting the `JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment "
-            "variable to `conservative`."
+        return
+    except Exception as initial_error:
+        # Check if this is a Julia registry error by looking at the error message
+        error_str = str(initial_error)
+        if (
+            "JuliaError" not in error_str
+            or "Unsatisfiable requirements detected" not in error_str
+        ):
+            raise initial_error
+
+        old_value = os.environ.get("JULIA_PKG_SERVER_REGISTRY_PREFERENCE", None)
+        if old_value == "eager":
+            raise initial_error
+
+        warnings.warn(
+            "Initial Julia registry operation failed. Attempting to use the `eager` registry flavor of the Julia "
+            "General registry from the Julia Pkg server (via the `JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment variable)."
         )
-        if old_value is not None:
-            warnings.warn(error_msg)
-            restore_juliaregistrypref(old_value)
-            raise e
-        else:
-            os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "conservative"
-            try:
-                f(args)
-            except Exception as e:
-                warnings.warn(error_msg)
-                restore_juliaregistrypref(old_value)
-                raise e
+        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
+        try:
+            f(*args)
+        finally:
+            if old_value is not None:
+                os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value
+            else:
+                del os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]

From 1fe1a7449f51028218280f2d8201e1eee9b4527e Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:10:07 +0000
Subject: [PATCH 08/14] refactor: simplify registry pref code

---
 pysr/julia_extensions.py       |  4 ++--
 pysr/julia_import.py           | 14 ++++++------
 pysr/julia_registry_helpers.py | 41 +++++-----------------------------
 3 files changed, 14 insertions(+), 45 deletions(-)

diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index b8fad157..950c292e 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -3,7 +3,7 @@
 from typing import Literal
 
 from .julia_import import Pkg, jl
-from .julia_registry_helpers import with_juliaregistrypref
+from .julia_registry_helpers import try_with_registry_fallback
 from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
 
 
@@ -53,7 +53,7 @@ def _add_package():
             Pkg.add(name=package_name, uuid=uuid_s)
             Pkg.resolve()
 
-        with_juliaregistrypref(_add_package)
+        try_with_registry_fallback(_add_package)
 
     # TODO: Protect against loading the same symbol from two packages,
     #       maybe with a @gensym here.
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index a1e5996e..e31e914e 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -4,12 +4,7 @@
 from types import ModuleType
 from typing import cast
 
-from .julia_registry_helpers import with_juliaregistrypref
-
-
-def import_juliacall():
-    import juliacall  # type: ignore
-
+from .julia_registry_helpers import try_with_registry_fallback
 
 # Check if JuliaCall is already loaded, and if so, warn the user
 # about the relevant environment variables. If not loaded,
@@ -49,7 +44,12 @@ def import_juliacall():
     # Deprecated; so just pass to juliacall
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
-with_juliaregistrypref(import_juliacall)
+
+def import_juliacall():
+    import juliacall  # type: ignore
+
+
+try_with_registry_fallback(import_juliacall)
 
 
 from juliacall import AnyValue  # type: ignore
diff --git a/pysr/julia_registry_helpers.py b/pysr/julia_registry_helpers.py
index fed82900..60f18a96 100644
--- a/pysr/julia_registry_helpers.py
+++ b/pysr/julia_registry_helpers.py
@@ -3,51 +3,20 @@
 import os
 import warnings
 from collections.abc import Callable
+from typing import TypeVar
 
+T = TypeVar("T")
 
-def backup_juliaregistrypref():
-    """Backup and potentially modify Julia registry preference.
 
-    Sets JULIA_PKG_SERVER_REGISTRY_PREFERENCE to 'eager' if not already set.
-    Returns the original value for later restoration.
-    """
-    old_value = os.environ.get("JULIA_PKG_SERVER_REGISTRY_PREFERENCE", None)
-    if old_value is None:
-        warnings.warn(
-            "Attempting to use the `eager` registry flavor of the Julia "
-            "General registry from the Julia Pkg server.\n"
-            "    If any errors are encountered, try setting the "
-            "`JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment variable to `conservative`."
-        )
-        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
-    return old_value
-
-
-def restore_juliaregistrypref(old_value: str | None):
-    """Restore the original Julia registry preference value."""
-    if old_value is None:
-        del os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]
-    else:
-        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value
-
-
-def with_juliaregistrypref(f: Callable[..., None], *args):
+def try_with_registry_fallback(f: Callable[..., T], *args, **kwargs) -> T:
     """Execute function with modified Julia registry preference.
 
     First tries with existing registry preference. If that fails with a Julia registry error,
     temporarily modifies the registry preference to 'eager'. Restores original preference after
     execution.
-
-    Parameters
-    ----------
-    f : Callable[..., None]
-        Function to execute. Should not return anything of importance.
-    *args : Any
-        Arguments to pass to the function.
     """
     try:
-        f(*args)
-        return
+        return f(*args, **kwargs)
     except Exception as initial_error:
         # Check if this is a Julia registry error by looking at the error message
         error_str = str(initial_error)
@@ -67,7 +36,7 @@ def with_juliaregistrypref(f: Callable[..., None], *args):
         )
         os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
         try:
-            f(*args)
+            return f(*args, **kwargs)
         finally:
             if old_value is not None:
                 os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value

From 30163aa8f2c3d54bdda34a44029403fee2f71f16 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:39:32 +0000
Subject: [PATCH 09/14] refactor: create global for preference key

---
 pysr/julia_registry_helpers.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/pysr/julia_registry_helpers.py b/pysr/julia_registry_helpers.py
index 60f18a96..fd99c212 100644
--- a/pysr/julia_registry_helpers.py
+++ b/pysr/julia_registry_helpers.py
@@ -7,6 +7,8 @@
 
 T = TypeVar("T")
 
+PREFERENCE_KEY = "JULIA_PKG_SERVER_REGISTRY_PREFERENCE"
+
 
 def try_with_registry_fallback(f: Callable[..., T], *args, **kwargs) -> T:
     """Execute function with modified Julia registry preference.
@@ -26,19 +28,19 @@ def try_with_registry_fallback(f: Callable[..., T], *args, **kwargs) -> T:
         ):
             raise initial_error
 
-        old_value = os.environ.get("JULIA_PKG_SERVER_REGISTRY_PREFERENCE", None)
+        old_value = os.environ.get(PREFERENCE_KEY, None)
         if old_value == "eager":
             raise initial_error
 
         warnings.warn(
             "Initial Julia registry operation failed. Attempting to use the `eager` registry flavor of the Julia "
-            "General registry from the Julia Pkg server (via the `JULIA_PKG_SERVER_REGISTRY_PREFERENCE` environment variable)."
+            + f"General registry from the Julia Pkg server (via the `{PREFERENCE_KEY}` environment variable)."
         )
-        os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = "eager"
+        os.environ[PREFERENCE_KEY] = "eager"
         try:
             return f(*args, **kwargs)
         finally:
             if old_value is not None:
-                os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"] = old_value
+                os.environ[PREFERENCE_KEY] = old_value
             else:
-                del os.environ["JULIA_PKG_SERVER_REGISTRY_PREFERENCE"]
+                del os.environ[PREFERENCE_KEY]

From 52b8e0abef87753a12be3940b183f172ed332c57 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:39:46 +0000
Subject: [PATCH 10/14] test: registry preference fallback

---
 pysr/test/test_startup.py | 77 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index a92010ce..a104b781 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -11,6 +11,7 @@
 
 from pysr import PySRRegressor
 from pysr.julia_import import jl_version
+from pysr.julia_registry_helpers import PREFERENCE_KEY, try_with_registry_fallback
 
 from .params import DEFAULT_NITERATIONS, DEFAULT_POPULATIONS
 
@@ -159,8 +160,82 @@ def test_notebook(self):
         self.assertEqual(result.returncode, 0)
 
 
+class TestRegistryHelper(unittest.TestCase):
+    """Test suite for Julia registry preference handling."""
+
+    def setUp(self):
+        """Store original environment."""
+        self.old_value = os.environ.get(PREFERENCE_KEY, None)
+
+    def tearDown(self):
+        """Restore original environment."""
+        if self.old_value is not None:
+            os.environ[PREFERENCE_KEY] = self.old_value
+        else:
+            os.environ.pop(PREFERENCE_KEY, None)
+
+    def test_successful_operation(self):
+        self.assertEqual(try_with_registry_fallback(lambda s: s, "success"), "success")
+
+    def test_non_julia_errors_reraised(self):
+        """Test that non-Julia errors are re-raised without modification."""
+        with self.assertRaises(SyntaxError) as context:
+            try_with_registry_fallback(lambda: exec("invalid syntax !@#$"))
+        self.assertNotIn("JuliaError", str(context.exception))
+
+    def test_julia_error_triggers_fallback(self):
+        """Test that Julia registry errors trigger the fallback to eager mode."""
+        os.environ[PREFERENCE_KEY] = "conservative"
+        recorded_env_vars = []
+
+        def failing_operation():
+            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
+            # Just fake the error message, as not sure how to realistically
+            # trigger a `conservative` registry-related error.
+            raise Exception(
+                "JuliaError\nUnsatisfiable requirements detected for package Test [8dfed614-e22c-5e08-85e1-65c5234f0b40]"
+            )
+
+        with self.assertWarns(Warning) as warn_context:
+            with self.assertRaises(Exception) as error_context:
+                try_with_registry_fallback(failing_operation)
+
+        self.assertIn("JuliaError", str(error_context.exception))
+        self.assertIn(
+            "Initial Julia registry operation failed. Attempting to use the `eager` registry flavor of the Julia",
+            str(warn_context.warning),
+        )
+
+        # Verify both modes are tried in order
+        self.assertEqual(recorded_env_vars, ["conservative", "eager"])
+
+        # Verify environment is restored
+        self.assertEqual(os.environ[PREFERENCE_KEY], "conservative")
+
+    def test_eager_mode_fails_directly(self):
+        """Test that eager mode errors don't trigger fallback."""
+        os.environ[PREFERENCE_KEY] = "eager"
+
+        recorded_env_vars = []
+        hits = 0
+
+        def failing_operation():
+            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
+            nonlocal hits
+            hits += 1
+            raise Exception(
+                "JuliaError\nUnsatisfiable requirements detected for package Test [8dfed614-e22c-5e08-85e1-65c5234f0b40]"
+            )
+
+        with self.assertRaises(Exception) as context:
+            try_with_registry_fallback(failing_operation)
+        self.assertIn("JuliaError", str(context.exception))
+        self.assertEqual(recorded_env_vars, ["eager"])  # Should only try eager mode
+        self.assertEqual(hits, 1)
+
+
 def runtests(just_tests=False):
-    tests = [TestStartup]
+    tests = [TestStartup, TestRegistryHelper]
     if just_tests:
         return tests
     suite = unittest.TestSuite()

From 400d22458bf622f93ec0cd9c395d7f7fbc59f351 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:40:40 +0000
Subject: [PATCH 11/14] refactor: private function name

---
 pysr/julia_import.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index e31e914e..4ea6b88d 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -45,11 +45,11 @@
     os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
 
 
-def import_juliacall():
+def _import_juliacall():
     import juliacall  # type: ignore
 
 
-try_with_registry_fallback(import_juliacall)
+try_with_registry_fallback(_import_juliacall)
 
 
 from juliacall import AnyValue  # type: ignore

From 25342995bda62b7a383d6607ecc591b930f6d2db Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:41:41 +0000
Subject: [PATCH 12/14] style: redundant comments

---
 pysr/test/test_startup.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index a104b781..61537a04 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -161,14 +161,12 @@ def test_notebook(self):
 
 
 class TestRegistryHelper(unittest.TestCase):
-    """Test suite for Julia registry preference handling."""
+    """Test the custom Julia registry preference handling."""
 
     def setUp(self):
-        """Store original environment."""
         self.old_value = os.environ.get(PREFERENCE_KEY, None)
 
     def tearDown(self):
-        """Restore original environment."""
         if self.old_value is not None:
             os.environ[PREFERENCE_KEY] = self.old_value
         else:
@@ -178,13 +176,11 @@ def test_successful_operation(self):
         self.assertEqual(try_with_registry_fallback(lambda s: s, "success"), "success")
 
     def test_non_julia_errors_reraised(self):
-        """Test that non-Julia errors are re-raised without modification."""
         with self.assertRaises(SyntaxError) as context:
             try_with_registry_fallback(lambda: exec("invalid syntax !@#$"))
         self.assertNotIn("JuliaError", str(context.exception))
 
     def test_julia_error_triggers_fallback(self):
-        """Test that Julia registry errors trigger the fallback to eager mode."""
         os.environ[PREFERENCE_KEY] = "conservative"
         recorded_env_vars = []
 
@@ -213,7 +209,6 @@ def failing_operation():
         self.assertEqual(os.environ[PREFERENCE_KEY], "conservative")
 
     def test_eager_mode_fails_directly(self):
-        """Test that eager mode errors don't trigger fallback."""
         os.environ[PREFERENCE_KEY] = "eager"
 
         recorded_env_vars = []

From 612d9ed39b501ad0335b26646c2a66dc2c240514 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 13:57:28 +0000
Subject: [PATCH 13/14] fix: fallback condition for registry switch

---
 pysr/julia_registry_helpers.py |  8 +++-----
 pysr/test/test_startup.py      | 13 ++++++-------
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/pysr/julia_registry_helpers.py b/pysr/julia_registry_helpers.py
index fd99c212..2c2162e7 100644
--- a/pysr/julia_registry_helpers.py
+++ b/pysr/julia_registry_helpers.py
@@ -21,11 +21,9 @@ def try_with_registry_fallback(f: Callable[..., T], *args, **kwargs) -> T:
         return f(*args, **kwargs)
     except Exception as initial_error:
         # Check if this is a Julia registry error by looking at the error message
-        error_str = str(initial_error)
-        if (
-            "JuliaError" not in error_str
-            or "Unsatisfiable requirements detected" not in error_str
-        ):
+        if "JuliaError" not in str(
+            type(initial_error)
+        ) or "Unsatisfiable requirements detected" not in str(initial_error):
             raise initial_error
 
         old_value = os.environ.get(PREFERENCE_KEY, None)
diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index 61537a04..5e2ce350 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -9,7 +9,7 @@
 
 import numpy as np
 
-from pysr import PySRRegressor
+from pysr import PySRRegressor, jl
 from pysr.julia_import import jl_version
 from pysr.julia_registry_helpers import PREFERENCE_KEY, try_with_registry_fallback
 
@@ -186,17 +186,16 @@ def test_julia_error_triggers_fallback(self):
 
         def failing_operation():
             recorded_env_vars.append(os.environ[PREFERENCE_KEY])
-            # Just fake the error message, as not sure how to realistically
-            # trigger a `conservative` registry-related error.
-            raise Exception(
-                "JuliaError\nUnsatisfiable requirements detected for package Test [8dfed614-e22c-5e08-85e1-65c5234f0b40]"
-            )
+            # Just add some package I know will not exist and also not be in the dependency chain:
+            jl.Pkg.add(name="AirspeedVelocity", version="100.0.0")
 
         with self.assertWarns(Warning) as warn_context:
             with self.assertRaises(Exception) as error_context:
                 try_with_registry_fallback(failing_operation)
 
-        self.assertIn("JuliaError", str(error_context.exception))
+        self.assertIn(
+            "Unsatisfiable requirements detected", str(error_context.exception)
+        )
         self.assertIn(
             "Initial Julia registry operation failed. Attempting to use the `eager` registry flavor of the Julia",
             str(warn_context.warning),

From 25c86395655086fa7bdfd8d05278854c81feea3c Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 14:08:55 +0000
Subject: [PATCH 14/14] test: make registry test more robust

---
 pysr/test/test_startup.py | 43 ++++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index 5e2ce350..4b2a450b 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -165,6 +165,16 @@ class TestRegistryHelper(unittest.TestCase):
 
     def setUp(self):
         self.old_value = os.environ.get(PREFERENCE_KEY, None)
+        self.recorded_env_vars = []
+        self.hits = 0
+
+        def failing_operation():
+            self.recorded_env_vars.append(os.environ[PREFERENCE_KEY])
+            self.hits += 1
+            # Just add some package I know will not exist and also not be in the dependency chain:
+            jl.Pkg.add(name="AirspeedVelocity", version="100.0.0")
+
+        self.failing_operation = failing_operation
 
     def tearDown(self):
         if self.old_value is not None:
@@ -182,16 +192,10 @@ def test_non_julia_errors_reraised(self):
 
     def test_julia_error_triggers_fallback(self):
         os.environ[PREFERENCE_KEY] = "conservative"
-        recorded_env_vars = []
-
-        def failing_operation():
-            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
-            # Just add some package I know will not exist and also not be in the dependency chain:
-            jl.Pkg.add(name="AirspeedVelocity", version="100.0.0")
 
         with self.assertWarns(Warning) as warn_context:
             with self.assertRaises(Exception) as error_context:
-                try_with_registry_fallback(failing_operation)
+                try_with_registry_fallback(self.failing_operation)
 
         self.assertIn(
             "Unsatisfiable requirements detected", str(error_context.exception)
@@ -202,7 +206,8 @@ def failing_operation():
         )
 
         # Verify both modes are tried in order
-        self.assertEqual(recorded_env_vars, ["conservative", "eager"])
+        self.assertEqual(self.recorded_env_vars, ["conservative", "eager"])
+        self.assertEqual(self.hits, 2)
 
         # Verify environment is restored
         self.assertEqual(os.environ[PREFERENCE_KEY], "conservative")
@@ -210,22 +215,14 @@ def failing_operation():
     def test_eager_mode_fails_directly(self):
         os.environ[PREFERENCE_KEY] = "eager"
 
-        recorded_env_vars = []
-        hits = 0
-
-        def failing_operation():
-            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
-            nonlocal hits
-            hits += 1
-            raise Exception(
-                "JuliaError\nUnsatisfiable requirements detected for package Test [8dfed614-e22c-5e08-85e1-65c5234f0b40]"
-            )
-
         with self.assertRaises(Exception) as context:
-            try_with_registry_fallback(failing_operation)
-        self.assertIn("JuliaError", str(context.exception))
-        self.assertEqual(recorded_env_vars, ["eager"])  # Should only try eager mode
-        self.assertEqual(hits, 1)
+            try_with_registry_fallback(self.failing_operation)
+
+        self.assertIn("Unsatisfiable requirements detected", str(context.exception))
+        self.assertEqual(
+            self.recorded_env_vars, ["eager"]
+        )  # Should only try eager mode
+        self.assertEqual(self.hits, 1)
 
 
 def runtests(just_tests=False):