diff --git a/dgf/src/plot/BUILD b/dgf/src/plot/BUILD index 102f070..1a07d28 100644 --- a/dgf/src/plot/BUILD +++ b/dgf/src/plot/BUILD @@ -27,6 +27,7 @@ py_library( srcs = ["pyvis.py"], deps = [ "//dgf/src/data:schema", + "//dgf/src/util:options", "//third_party/py/pyvis", ], ) diff --git a/dgf/src/plot/pyvis.py b/dgf/src/plot/pyvis.py index b4eae98..9a33be2 100644 --- a/dgf/src/plot/pyvis.py +++ b/dgf/src/plot/pyvis.py @@ -12,13 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Plotting of graph elements using the pyvis library.""" +"""Plotting utilities for graph schemas using pyvis.""" -from typing import Any, Dict +from typing import Any from dgf.src.data import schema as schema_lib +from dgf.src.util import options as options_lib from pyvis import network as pyvis_network +Context = options_lib.Context +options = options_lib.Manager("pyvis_options") +default_options = options.default_options +option_context = options.context + def _html_label(name: str, features: list[str]) -> str: """Creates a node/edge title for hover information.""" @@ -31,7 +37,7 @@ def plot_schema( schema: schema_lib.GraphSchema, features: bool = True, *, - pyvis_kwargs: Dict[Any, Any] = {}, + pyvis_kwargs: dict[Any, Any] | None = None, ) -> pyvis_network.Network: """Plots the graph schema's meta-graph (i.e., its nodesets and edgesets). @@ -39,12 +45,20 @@ def plot_schema( schema: The `GraphSchema` object to plot. features: If true, display the node and edges features in the title (hover). pyvis_kwargs: Additional keyword arguments to pass to the - `pyvis.network.Network` constructor. + `pyvis.network.Network` constructor. These override any default or + contextual pyvis options. Returns: A `pyvis.network.Network` object representing the graph schema. """ - net = pyvis_network.Network(directed=True, **pyvis_kwargs) + final_kwargs: dict[str, Any] = {"directed": True} + final_kwargs.update(options.to_dict()) + + if pyvis_kwargs is not None: + final_kwargs.update(pyvis_kwargs) + + clean_kwargs = {k: v for k, v in final_kwargs.items() if v is not None} + net = pyvis_network.Network(**clean_kwargs) # Add nodes for node_set_name in sorted(schema.node_sets.keys()): diff --git a/dgf/src/plot/pyvis_test.py b/dgf/src/plot/pyvis_test.py index dcba57f..e857d28 100644 --- a/dgf/src/plot/pyvis_test.py +++ b/dgf/src/plot/pyvis_test.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for dgf.src.plot.pyvis.""" - from absl.testing import absltest from dgf.src.plot import pyvis as pyvis_plot from dgf.src.util import gen_test_graph @@ -21,6 +19,14 @@ class PyvisTest(absltest.TestCase): + def setUp(self): + super().setUp() + pyvis_plot.options.reset() + + def tearDown(self): + pyvis_plot.options.reset() + super().tearDown() + def test_plot_schema(self): schema = gen_test_graph.generate_schema() net = pyvis_plot.plot_schema(schema, features=True) @@ -34,6 +40,31 @@ def test_plot_schema(self): self.assertIn("n1", html) self.assertIn("n2", html) + def test_plot_schema_with_none_option_fallback(self): + schema = gen_test_graph.generate_schema() + pyvis_plot.options.set("height", "500px") + + # Setting height=None instructs plot_schema to omit the kwarg + # (falling back to pyvis default 600px). + with pyvis_plot.option_context(height=None): + net = pyvis_plot.plot_schema(schema) + self.assertEqual(net.height, "600px") + + net = pyvis_plot.plot_schema(schema) + self.assertEqual(net.height, "500px") # Custom height is active again. + + def test_plot_schema_with_default_options_and_override(self): + schema = gen_test_graph.generate_schema() + pyvis_plot.options.set("height", "500px") + + net1 = pyvis_plot.plot_schema(schema) + self.assertEqual(net1.height, "500px") # Default options applied. + + net2 = pyvis_plot.plot_schema(schema, pyvis_kwargs={"height": "750px"}) + self.assertEqual(net2.height, "750px") # Call-specific kwargs override. + + self.assertEqual(pyvis_plot.options.get("height"), "500px") # Unchanged. + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/util/BUILD b/dgf/src/util/BUILD index 82d1b54..48b3346 100644 --- a/dgf/src/util/BUILD +++ b/dgf/src/util/BUILD @@ -387,3 +387,17 @@ py_test( # numpy dep, ], ) + +py_library( + name = "options", + srcs = ["options.py"], +) + +py_test( + name = "options_test", + srcs = ["options_test.py"], + deps = [ + ":options", + # absl/testing:absltest dep, + ], +) diff --git a/dgf/src/util/options.py b/dgf/src/util/options.py new file mode 100644 index 0000000..e13fa14 --- /dev/null +++ b/dgf/src/util/options.py @@ -0,0 +1,160 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Scoped, hierarchical options and configuration context managers.""" + +from collections.abc import Iterator, Mapping, MutableMapping +import contextlib +import contextvars +from typing import Any + +_SENTINEL = object() + + +class Context(Mapping[str, Any]): + """Immutable context scope holding a dictionary of options and an optional parent scope.""" + + def __init__( + self, + options: Mapping[str, Any] | None = None, + parent: "Context | None" = None, + ): + self._options: dict[str, Any] = dict(options) if options is not None else {} + self._parent = parent + + def get_option(self, key: str, default: Any = None) -> Any: + """Gets the option value from this context or its parent scopes.""" + if key in self._options: + return self._options[key] + if self._parent is not None: + return self._parent.get_option(key, default) + return default + + def get(self, key: str, default: Any = None) -> Any: + """Gets the option value from this context or its parent scopes.""" + return self.get_option(key, default) + + def __getitem__(self, key: str) -> Any: + val = self.get_option(key, default=_SENTINEL) + if val is _SENTINEL: + raise KeyError(key) + return val + + def __contains__(self, key: object) -> bool: + if key in self._options: + return True + return self._parent is not None and key in self._parent + + def __iter__(self) -> Iterator[str]: + return iter(self.to_dict()) + + def __len__(self) -> int: + return len(self.to_dict()) + + def with_option(self, key: str, value: Any) -> "Context": + """Returns a new Context with the option set (Copy-on-Write).""" + new_options = self._options.copy() + new_options[key] = value + return Context(new_options, parent=self._parent) + + def with_options(self, **kwargs: Any) -> "Context": + """Returns a new Context with the given options updated (Copy-on-Write).""" + new_options = self._options.copy() + new_options.update(kwargs) + return Context(new_options, parent=self._parent) + + def without_option(self, key: str | None = None) -> "Context": + """Returns a new Context with the option(s) reset (Copy-on-Write).""" + if key is None: + new_options = {} + else: + new_options = self._options.copy() + new_options.pop(key, None) + return Context(new_options, parent=self._parent) + + def to_dict(self) -> dict[str, Any]: + """Returns the merged options dictionary up to this context level.""" + merged = self._parent.to_dict() if self._parent is not None else {} + merged.update(self._options) + return merged + + +class Manager(MutableMapping[str, Any]): + """Thread-safe and coroutine-safe options manager for a subsystem.""" + + def __init__( + self, + name: str, + defaults: Mapping[str, Any] | None = None, + ): + self._default_options = Context(defaults) + self._current_context: contextvars.ContextVar[Context] = ( + contextvars.ContextVar(name, default=self._default_options) + ) + + @property + def default_options(self) -> Context: + """Returns the base/default options context.""" + return self._default_options + + def get(self, key: str, default: Any = None) -> Any: + """Gets the active option value from the current context stack.""" + return self._current_context.get().get_option(key, default) + + def set(self, key: str, value: Any) -> None: + """Sets an option in the active context scope using Copy-on-Write.""" + current_context = self._current_context.get() + self._current_context.set(current_context.with_option(key, value)) + + def reset(self, key: str | None = None) -> None: + """Resets an option or all options in the active context scope using Copy-on-Write.""" + self._current_context.set(self._current_context.get().without_option(key)) + + def __getitem__(self, key: str) -> Any: + return self._current_context.get()[key] + + def __setitem__(self, key: str, value: Any) -> None: + self.set(key, value) + + def __delitem__(self, key: str) -> None: + if key not in self: + raise KeyError(key) + self.reset(key) + + def __contains__(self, key: object) -> bool: + return key in self._current_context.get() + + def __iter__(self) -> Iterator[str]: + return iter(self._current_context.get()) + + def __len__(self) -> int: + return len(self._current_context.get()) + + def to_dict(self) -> dict[str, Any]: + """Returns the merged dictionary of active options.""" + return self._current_context.get().to_dict() + + @contextlib.contextmanager + def context(self, **kwargs: Any) -> Iterator[Context]: + """Context manager to temporarily override options in a stack-safe scope.""" + ctx = Context(kwargs.copy(), parent=self._current_context.get()) + token = self._current_context.set(ctx) + try: + yield ctx + finally: + self._current_context.reset(token) + + def __call__(self, **kwargs: Any) -> Any: + """Allows manager instance to be used directly as a context manager: `with options(...)`.""" + return self.context(**kwargs) diff --git a/dgf/src/util/options_test.py b/dgf/src/util/options_test.py new file mode 100644 index 0000000..e4588f8 --- /dev/null +++ b/dgf/src/util/options_test.py @@ -0,0 +1,175 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Mapping, MutableMapping +import concurrent.futures +import time + +from absl.testing import absltest +from dgf.src.util import options as options_lib + + +class OptionsTest(absltest.TestCase): + + def test_context_basic(self): + ctx = options_lib.Context({"a": 1, "b": 2}) + self.assertEqual(ctx.get_option("a"), 1) + self.assertEqual(ctx.get_option("b"), 2) + self.assertEqual(ctx.get("a"), 1) + self.assertEqual(ctx["a"], 1) + self.assertIsNone(ctx.get_option("c")) + self.assertEqual(ctx.get_option("c", default=10), 10) + self.assertEqual(ctx.get("c", 10), 10) + self.assertIn("a", ctx) + self.assertNotIn("c", ctx) + self.assertLen(ctx, 2) + self.assertEqual(set(ctx), {"a", "b"}) + self.assertEqual(ctx.to_dict(), {"a": 1, "b": 2}) + + with self.assertRaises(KeyError): + _ = ctx["c"] + + self.assertIsInstance(ctx, Mapping) + self.assertNotIsInstance(ctx, MutableMapping) + self.assertFalse(hasattr(ctx, "set_option")) + self.assertFalse(hasattr(ctx, "reset_option")) + self.assertFalse(hasattr(ctx, "__setitem__")) + + derived = options_lib.Context(ctx) + self.assertEqual(derived.to_dict(), {"a": 1, "b": 2}) + + def test_context_none_values(self): + ctx = options_lib.Context({"seed": None}) + self.assertIn("seed", ctx) + self.assertIsNone(ctx.get_option("seed", default=42)) + self.assertIsNone(ctx["seed"]) + + parent = options_lib.Context({"seed": 42}) + child = options_lib.Context({"seed": None}, parent=parent) + self.assertIn("seed", child) + self.assertIsNone(child.get_option("seed", default=99)) + self.assertIsNone(child["seed"]) + + def test_defensive_dict_copying(self): + orig = {"a": 1, "b": 2} + ctx = options_lib.Context(orig) + orig["a"] = 999 + orig["c"] = 3 + self.assertEqual(ctx.get_option("a"), 1) + self.assertNotIn("c", ctx) + + def test_context_hierarchy(self): + parent = options_lib.Context({"a": 1, "b": 2}) + child = options_lib.Context({"b": 20, "c": 30}, parent=parent) + + self.assertEqual(child.get_option("a"), 1) # From parent + self.assertEqual(child.get_option("b"), 20) # Overridden in child + self.assertEqual(child.get_option("c"), 30) # From child + self.assertEqual(child["a"], 1) + + self.assertEqual(child.to_dict(), {"a": 1, "b": 20, "c": 30}) + + def test_context_copy_on_write_derivation(self): + ctx1 = options_lib.Context({"a": 1, "b": 2}) + + ctx2 = ctx1.with_option("b", 20).with_option("c", 30) + self.assertEqual(ctx1.to_dict(), {"a": 1, "b": 2}) # ctx1 untouched + self.assertEqual(ctx2.to_dict(), {"a": 1, "b": 20, "c": 30}) + + ctx3 = ctx1.with_options(b=200, d=400) + self.assertEqual(ctx1.to_dict(), {"a": 1, "b": 2}) # ctx1 untouched + self.assertEqual(ctx3.to_dict(), {"a": 1, "b": 200, "d": 400}) + + ctx4 = ctx2.without_option("b") + self.assertEqual(ctx4.to_dict(), {"a": 1, "c": 30}) + + ctx5 = ctx2.without_option() + self.assertEqual(ctx5.to_dict(), {}) + + def test_manager_context(self): + manager = options_lib.Manager("test_manager", defaults={"width": 100}) + self.assertEqual(manager.get("width"), 100) + self.assertIsNone(manager.get("height")) + + with manager.context(width=200, height=50) as ctx: + self.assertEqual(manager.get("width"), 200) + self.assertEqual(manager.get("height"), 50) + self.assertEqual(ctx.get_option("width"), 200) + self.assertEqual(manager.to_dict(), {"width": 200, "height": 50}) + + with manager(height=80, color="blue", layout=None): + self.assertEqual(manager.get("width"), 200) # Inherited from outer + self.assertEqual(manager.get("height"), 80) # Overridden in inner + self.assertEqual(manager.get("color"), "blue") + self.assertIsNone(manager.get("layout", default="force")) + + self.assertEqual(manager.get("height"), 50) + self.assertIsNone(manager.get("color")) + + self.assertEqual(manager.get("width"), 100) + self.assertIsNone(manager.get("height")) + + def test_manager_mapping_protocol(self): + manager = options_lib.Manager("test_mapping", defaults={"a": 1, "b": 2}) + self.assertIsInstance(manager, MutableMapping) + self.assertLen(manager, 2) + self.assertEqual(set(manager), {"a", "b"}) + self.assertIn("a", manager) + self.assertNotIn("c", manager) + self.assertEqual(manager["a"], 1) + + manager["c"] = 3 + self.assertLen(manager, 3) + self.assertEqual(manager["c"], 3) + self.assertIn("c", manager) + + del manager["c"] + self.assertNotIn("c", manager) + self.assertLen(manager, 2) + + with self.assertRaises(KeyError): + del manager["non_existent"] + + with self.assertRaises(KeyError): + _ = manager["missing"] + + with manager(a=10, d=40): + self.assertLen(manager, 3) + self.assertEqual(set(manager), {"a", "b", "d"}) + self.assertEqual(manager["a"], 10) + self.assertEqual(manager["d"], 40) + + self.assertLen(manager, 2) + self.assertEqual(manager["a"], 1) + self.assertNotIn("d", manager) + + def test_cow_thread_isolation(self): + manager = options_lib.Manager("thread_cow", defaults={"counter": 0}) + + def worker(val: int) -> int: + manager["counter"] = val + time.sleep(0.01) + return manager["counter"] + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: + f1 = executor.submit(worker, 100) + f2 = executor.submit(worker, 200) + self.assertEqual(f1.result(), 100) + self.assertEqual(f2.result(), 200) + + self.assertEqual(manager["counter"], 0) + + +if __name__ == "__main__": + absltest.main()