From acf502987f3279e0638da904ea099c5cbcbea21c Mon Sep 17 00:00:00 2001
From: Erik Welch <erik.n.welch@gmail.com>
Date: Mon, 7 Oct 2024 09:53:07 -0500
Subject: [PATCH] `nx-cugraph`: add `NX_CUGRAPH_AUTOCONFIG=True` env var to
 enable full zero-code change (#4685)

This is for convenience and sets or updates `NETWORKX` environment variables.

Do we like `NX_CUGRAPH` as the env var name? What should we consider a true value: "True" (case-insensitive) or a non-empty value?

This works with the latest dev version of NetworkX. I have not yet tried it with older NetworkX versions.

Authors:
  - Erik Welch (https://github.com/eriknw)
  - https://github.com/jakirkham

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - Ralph Liu (https://github.com/nv-rliu)

URL: https://github.com/rapidsai/cugraph/pull/4685
---
 python/nx-cugraph/_nx_cugraph/__init__.py | 39 +++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
index a5e45979fe2..fc0bea47180 100644
--- a/python/nx-cugraph/_nx_cugraph/__init__.py
+++ b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -301,6 +301,45 @@ def get_info():
         .lower()
         == "true",
     }
+
+    # Enable zero-code change usage with a simple environment variable
+    # by setting or updating other NETWORKX environment variables.
+    if os.environ.get("NX_CUGRAPH_AUTOCONFIG", "").strip().lower() == "true":
+        from itertools import chain
+
+        def update_env_var(varname):
+            """Add "cugraph" to a list of backend names environment variable."""
+            if varname not in os.environ:
+                os.environ[varname] = "cugraph"
+                return
+            string = os.environ[varname]
+            vals = [
+                stripped for x in string.strip().split(",") if (stripped := x.strip())
+            ]
+            if "cugraph" not in vals:
+                # Should we append or prepend? Let's be first!
+                os.environ[varname] = ",".join(chain(["cugraph"], vals))
+
+        # Automatically convert NetworkX Graphs to nx-cugraph for algorithms
+        if (varname := "NETWORKX_BACKEND_PRIORITY_ALGOS") in os.environ:
+            # "*_ALGOS" is given priority in NetworkX >=3.4
+            update_env_var(varname)
+            # But update this too to "just work" if users mix env vars and nx versions
+            os.environ["NETWORKX_BACKEND_PRIORITY"] = os.environ[varname]
+        else:
+            update_env_var("NETWORKX_BACKEND_PRIORITY")
+        # And for older NetworkX versions
+        update_env_var("NETWORKX_AUTOMATIC_BACKENDS")  # For NetworkX 3.2
+        update_env_var("NETWORKX_GRAPH_CONVERT")  # For NetworkX 3.0 and 3.1
+        # Automatically create nx-cugraph Graph from graph generators
+        update_env_var("NETWORKX_BACKEND_PRIORITY_GENERATORS")
+        # Run default NetworkX implementation (in >=3.4) if not implemented by nx-cugraph
+        if (varname := "NETWORKX_FALLBACK_TO_NX") not in os.environ:
+            os.environ[varname] = "true"
+        # Cache graph conversions (default is False in NetworkX 3.2
+        if (varname := "NETWORKX_CACHE_CONVERTED_GRAPHS") not in os.environ:
+            os.environ[varname] = "true"
+
     return d