src/dctv/lazygraph.py

# Copyright (C) 2020 The Android Open Source Project
#
# 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
#
#      http:#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.
"""NetworkX DiGraph that expands nodes as needed"""

from contextlib import contextmanager
import networkx as nx

class LazyNodeDict(dict):
  """Replacement dictionary for LazyDiGraph implementation"""

  def __init__(self, expand, expanded):
    dict.__init__(self)
    self._expand = expand
    self._expanded = expanded

  def __getitem__(self, node):
    if node not in self._expanded:
      self._expand(node)
    return super().__getitem__(node)

  __slots__ = "_expand", "_expanded"

class LazyDiGraph(nx.DiGraph):
  """Class that lazily expands its nodes using a callback"""

  def __init__(self, expand):
    # Init expansion hook before parent init
    self.__expanded = set()
    self.__expand_function = expand
    self.node_dict_factory = self.__make_lazy_node_dict
    super().__init__()
    self.__expansion_disabled = False

  def __make_lazy_node_dict(self):
    return LazyNodeDict(self.__do_expand, self.__expanded)

  def __do_expand(self, node):
    if not self.__expansion_disabled:
      self.__expanded.add(node)
      self.__expand_function(self, node)

  @contextmanager
  def expansion_disabled(self):
    """Context manager that temporarily disables expansion"""
    old_disabled_expansion = self.__expansion_disabled
    self.__expansion_disabled = True
    try:
      yield
    finally:
      self.__expansion_disabled = old_disabled_expansion