Skip to content

Commit

Permalink
Add df.nx.set_properties(**kwargs)
Browse files Browse the repository at this point in the history
  • Loading branch information
eriknw committed Jul 23, 2024
1 parent 4896933 commit 6e56e9f
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 0 deletions.
48 changes: 48 additions & 0 deletions nx_pandas/_patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,54 @@ def __dir__(self):
attrs.remove("edge_key")
return attrs

def set_properties(
self,
*,
source=None,
target=None,
edge_key=None,
is_directed=None,
is_multigraph=None,
cache_enabled=None,
):
"""Set many graph properties (i.e., ``df.nx`` attributes) at once.
Return the original DataFrame to allow method chaining. For example::
>>> df = pd.read_csv("my_data.csv").nx.set_properties(is_directed=False)
This is a bulk transaction, so either all given attributes will be updated,
or nothing will be set if there was an exception.
"""
prev = {}
cur = {}
if source is not None:
prev["_source"] = self._source
cur["source"] = source
if target is not None:
prev["_target"] = self._target
cur["target"] = target
if is_directed is not None:
prev["is_directed"] = self.is_directed
cur["is_directed"] = is_directed
if is_multigraph is not None:
prev["is_multigraph"] = self.is_multigraph
cur["is_multigraph"] = is_multigraph
if edge_key is not None:
prev["_edge_key"] = self._edge_key
cur["edge_key"] = edge_key
if cache_enabled is not None:
prev["cache_enabled"] = self.cache_enabled
cur["cache_enabled"] = cache_enabled
try:
for attr, val in cur.items():
setattr(self, attr, val)
except Exception:
for attr, val in prev.items():
setattr(self, attr, val)
raise
return self._df


def _attr_raise_if_invalid_graph(df, attr):
try:
Expand Down
36 changes: 36 additions & 0 deletions nx_pandas/tests/test_df_nx_attrs.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,3 +110,39 @@ def test_set_attrs(df):
df.nx.edge_key = None
with pytest.raises(AttributeError, match="to be used as a networkx graph"):
df.__networkx_backend__


def test_set_properties(df):
df2 = df.nx.set_properties(
source="target",
target="source",
edge_key="foo",
is_directed=False,
is_multigraph=True,
cache_enabled=True,
)
assert df is df2
assert df.nx.source == "target"
assert df.nx.target == "source"
assert df.nx.edge_key == "foo"
assert df.nx.is_directed is False
assert df.nx.is_multigraph is True
assert df.nx.cache_enabled is True
with pytest.raises(
AttributeError, match="'edge_key' attribute only exists for multigraphs"
):
df.nx.set_properties(
source="source",
target="target",
is_directed=True,
is_multigraph=False,
cache_enabled=False,
edge_key="BAD",
)
# Unchanged
assert df.nx.source == "target"
assert df.nx.target == "source"
assert df.nx.edge_key == "foo"
assert df.nx.is_directed is False
assert df.nx.is_multigraph is True
assert df.nx.cache_enabled is True

0 comments on commit 6e56e9f

Please sign in to comment.