Format with black

tomalrussell · Mar 23, 2022 · 92995c2 · 92995c2
1 parent 4cd2d8a
commit 92995c2
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 33 deletions.
diff --git a/setup.py b/setup.py
@@ -54,7 +54,7 @@ def readme():
         # eg:
         #   'rst': ['docutils>=0.11'],
         #   ':python_version=="2.6"': ['argparse'],
-        "networkx": ['networkx>=1.9']
+        "networkx": ["networkx>=1.9"],
     },
     entry_points={
         "console_scripts": [

diff --git a/src/snkit/network.py b/src/snkit/network.py
@@ -10,8 +10,10 @@
 import numpy as np
 import pandas
 import shapely.errors
+
 try:
     import networkx as nx
+
     USE_NX = True
 except ImportError:
     USE_NX = False
@@ -680,11 +682,10 @@ def set_precision(geom, precision):
     return shape(geom_mapping)
 
 
-def to_networkx(network,directed=False,weight_col=None):
-    """Return a networkx graph
-    """
+def to_networkx(network, directed=False, weight_col=None):
+    """Return a networkx graph"""
     if not USE_NX:
-        raise ImportError('No module named networkx')
+        raise ImportError("No module named networkx")
     else:
         # init graph
         if not directed:
@@ -694,44 +695,57 @@ def to_networkx(network,directed=False,weight_col=None):
         # get nodes from network data
         G.add_nodes_from(network.nodes.id.to_list())
         # add nodal positions from geom
-        network.nodes['pos'] = list(zip(network.nodes.geometry.x, network.nodes.geometry.y))
-        pos = network.nodes.set_index('id').to_dict()['pos']
-        for n,p in pos.items():
-            G.nodes[n]['pos'] = p
+        network.nodes["pos"] = list(
+            zip(network.nodes.geometry.x, network.nodes.geometry.y)
+        )
+        pos = network.nodes.set_index("id").to_dict()["pos"]
+        for n, p in pos.items():
+            G.nodes[n]["pos"] = p
         # get edges from network data
         if weight_col is None:
-            network.edges['weight'] = network.edges.geometry.length
-            edges_as_list = list(zip(network.edges.from_id,network.edges.to_id,network.edges.weight))
+            network.edges["weight"] = network.edges.geometry.length
+            edges_as_list = list(
+                zip(network.edges.from_id, network.edges.to_id, network.edges.weight)
+            )
         else:
-            edges_as_list = list(zip(network.edges.from_id,network.edges.to_id,network.edges[weight_col]))
+            edges_as_list = list(
+                zip(
+                    network.edges.from_id,
+                    network.edges.to_id,
+                    network.edges[weight_col],
+                )
+            )
         # add edges to graph
         G.add_weighted_edges_from(edges_as_list)
         return G
 
 
 def get_connected_components(network):
-    """Get connected components within network and id to each individual graph
-    """
+    """Get connected components within network and id to each individual graph"""
     if not USE_NX:
-        raise ImportError('No module named networkx')
+        raise ImportError("No module named networkx")
     else:
         G = to_networkx(network)
-        return sorted(nx.connected_components(G), key = len, reverse=True)
+        return sorted(nx.connected_components(G), key=len, reverse=True)
 
 
-def add_component_ids(network,id_col='component_id'):
-    """Add column of component IDs to network data
-    """
+def add_component_ids(network, id_col="component_id"):
+    """Add column of component IDs to network data"""
     # get connected components
     connected_parts = get_connected_components(network)
     # add unique id to each graph
-    network.edges[id_col] = 0 # init id_col
-    network.nodes[id_col] = 0 # init id_col
+    network.edges[id_col] = 0  # init id_col
+    network.nodes[id_col] = 0  # init id_col
     for count, part in enumerate(connected_parts):
         # edges
-        network.edges.loc[ (network.edges.from_id.isin(list(part))) | \
-                           (network.edges.to_id.isin(list(part))), id_col ] = count + 1
+        network.edges.loc[
+            (network.edges.from_id.isin(list(part)))
+            | (network.edges.to_id.isin(list(part))),
+            id_col,
+        ] = (
+            count + 1
+        )
         # nodes
-        network.nodes.loc[ (network.nodes.id.isin(list(part))), id_col] = count + 1 
+        network.nodes.loc[(network.nodes.id.isin(list(part))), id_col] = count + 1
     # return
     return network
diff --git a/tests/test_init.py b/tests/test_init.py
@@ -14,6 +14,7 @@
 try:
     import networkx as nx
     from networkx.utils.misc import graphs_equal
+
     USE_NX = True
 except ImportError:
     USE_NX = False
@@ -361,15 +362,14 @@ def test_passing_slice():
 
 @mark.skipif(not USE_NX, reason="networkx not available")
 def test_to_networkx(connected):
-    '''test to networkx
-    '''
-    connected.nodes['id'] = ['n'+str(i) for i in connected.nodes.index]
+    """Test conversion to networkx"""
+    connected.nodes["id"] = ["n" + str(i) for i in connected.nodes.index]
     connected = snkit.network.add_topology(connected)
     G = snkit.network.to_networkx(connected)
-    
+
     G_true = nx.Graph()
-    G_true.add_node('n0',pos=(0, 0))
-    G_true.add_node('n1',pos=(0, 2))
-    G_true.add_edge('n0', 'n1', weight=2)
-    
-    assert graphs_equal(G, G_true)
+    G_true.add_node("n0", pos=(0, 0))
+    G_true.add_node("n1", pos=(0, 2))
+    G_true.add_edge("n0", "n1", weight=2)
+
+    assert graphs_equal(G, G_true)