Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for automatic dict and set attributes in module io #1075

Merged
merged 13 commits into from
Oct 24, 2023
Merged
16 changes: 14 additions & 2 deletions osmnx/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -421,6 +421,16 @@ def _convert_node_attr_types(G, dtypes=None):
G : networkx.MultiDiGraph
"""
for _, data in G.nodes(data=True):
# first, eval stringified lists, dicts or sets to convert them to objects
# node attributes might have a single value, or a list if simplified
for attr, value in data.items():
# check for stringified lists
if (value.startswith("[") and value.endswith("]")) or (
value.startswith("{") and value.endswith("}")
):
with contextlib.suppress(SyntaxError, ValueError):
data[attr] = ast.literal_eval(value)

for attr in data.keys() & dtypes.keys():
data[attr] = dtypes[attr](data[attr])
return G
Expand All @@ -446,10 +456,12 @@ def _convert_edge_attr_types(G, dtypes=None):
# remove extraneous "id" attribute added by graphml saving
data.pop("id", None)

# first, eval stringified lists to convert them to list objects
# first, eval stringified lists, dicts or sets to convert them to objects
# edge attributes might have a single value, or a list if simplified
for attr, value in data.items():
if value.startswith("[") and value.endswith("]"):
if (value.startswith("[") and value.endswith("]")) or (
value.startswith("{") and value.endswith("}")
):
with contextlib.suppress(SyntaxError, ValueError):
data[attr] = ast.literal_eval(value)

Expand Down
19 changes: 19 additions & 0 deletions tests/test_osmnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -474,6 +474,25 @@ def test_graph_save_load():
edge_attrs = {n: bool(b) for n, b in zip(G.edges, bools)}
nx.set_edge_attributes(G, edge_attrs, attr_name)

# create list, set, and dict attributes for nodes and edges
rand_ints_nodes = np.random.randint(0, 10, len(G.nodes))
rand_ints_edges = np.random.randint(0, 10, len(G.edges))
list_name = "test_list"
list_node_attrs = {n: [n, r] for n, r in zip(G.nodes, rand_ints_nodes)}
nx.set_node_attributes(G, list_node_attrs, list_name)
list_edge_attrs = {e: [e, r] for e, r in zip(G.edges, rand_ints_edges)}
nx.set_edge_attributes(G, list_edge_attrs, list_name)
set_name = "test_set"
set_node_attrs = {n: {n, r} for n, r in zip(G.nodes, rand_ints_nodes)}
nx.set_node_attributes(G, set_node_attrs, set_name)
set_edge_attrs = {e: {e, r} for e, r in zip(G.edges, rand_ints_edges)}
nx.set_edge_attributes(G, set_edge_attrs, set_name)
dict_name = "test_dict"
dict_node_attrs = {n: {n: r} for n, r in zip(G.nodes, rand_ints_nodes)}
nx.set_node_attributes(G, dict_node_attrs, dict_name)
dict_edge_attrs = {e: {e: r} for e, r in zip(G.edges, rand_ints_edges)}
nx.set_edge_attributes(G, dict_edge_attrs, dict_name)

# save/load graph as graphml file
ox.save_graphml(G, gephi=True)
ox.save_graphml(G, gephi=False)
Expand Down