Change format of function info in dict

rapidsai · Sep 27, 2023 · 7950c4e · 7950c4e
1 parent 67e3cc6
commit 7950c4e
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 30 deletions.
diff --git a/python/nx-cugraph/_nx_cugraph/__init__.py b/python/nx-cugraph/_nx_cugraph/__init__.py
@@ -19,42 +19,53 @@
 $ python _nx_cugraph/__init__.py
 """
 
+# Entries between BEGIN and END are automatically generated
+_info = {
+    "backend_name": "cugraph",
+    "project": "nx-cugraph",
+    "package": "nx_cugraph",
+    "url": "https://github.com/rapidsai/cugraph/tree/branch-23.10/python/nx-cugraph",
+    "short_summary": "GPU-accelerated backend.",
+    # "description": "TODO",
+    "functions": {
+        # BEGIN: functions
+        "betweenness_centrality",
+        "edge_betweenness_centrality",
+        "louvain_communities",
+        # END: functions
+    },
+    "extra_docstrings": {
+        # BEGIN: extra_docstrings
+        "betweenness_centrality": "`weight` parameter is not yet supported.",
+        "edge_betweenness_centrality": "`weight` parameter is not yet supported.",
+        "louvain_communities": "`threshold` and `seed` parameters are currently ignored.",
+        # END: extra_docstrings
+    },
+    "extra_parameters": {
+        # BEGIN: extra_parameters
+        "louvain_communities": {
+            "max_level : int, optional": "Upper limit of the number of macro-iterations.",
+        },
+        # END: extra_parameters
+    },
+}
+
 
 def get_info():
     """Target of ``networkx.plugin_info`` entry point.
 
     This tells NetworkX about the cugraph backend without importing nx_cugraph.
     """
-    # Entries between BEGIN and END are automatically generated
-    return {
-        "backend_name": "cugraph",
-        "project": "nx-cugraph",
-        "package": "nx_cugraph",
-        "url": "https://github.com/rapidsai/cugraph/tree/branch-23.10/python/nx-cugraph",
-        "short_summary": "GPU-accelerated backend.",
-        # "description": "TODO",
-        "functions": {
-            # BEGIN: functions
-            "betweenness_centrality",
-            "edge_betweenness_centrality",
-            "louvain_communities",
-            # END: functions
-        },
-        "extra_docstrings": {
-            # BEGIN: extra_docstrings
-            "betweenness_centrality": "`weight` parameter is not yet supported.",
-            "edge_betweenness_centrality": "`weight` parameter is not yet supported.",
-            "louvain_communities": "`threshold` and `seed` parameters are currently ignored.",
-            # END: extra_docstrings
-        },
-        "extra_parameters": {
-            # BEGIN: extra_parameters
-            "louvain_communities": {
-                "max_level : int, optional": "Upper limit of the number of macro-iterations.",
-            },
-            # END: extra_parameters
-        },
+    # Convert to e.g. `{"functions" {"myfunc": {"extra_docstrings": ...}}}`
+    d = _info.copy()
+    func_keys = ("extra_docstrings", "extra_parameters")
+    d["functions"] = {
+        func: {key: vals[func] for key in func_keys if func in (vals := d[key])}
+        for func in d["functions"]
     }
+    for key in func_keys:
+        del d[key]
+    return d
 
 
 __version__ = "23.10.00"

diff --git a/python/nx-cugraph/_nx_cugraph/core.py b/python/nx-cugraph/_nx_cugraph/core.py
@@ -36,7 +36,7 @@ def get_extra_parameters(functions=None):
     return {key: val.extra_params for key, val in functions.items() if val.extra_params}
 
 
-def update_text(text, lines_to_add, target, indent=" " * 12):
+def update_text(text, lines_to_add, target, indent=" " * 8):
     begin = f"# BEGIN: {target}\n"
     end = f"# END: {target}\n"
     start = text.index(begin)