fix lint

Author: bjchambers

libs/community/langchain_community/graph_vectorstores/visualize.py

@@ -1,5 +1,5 @@
-from typing import TYPE_CHECKING, Dict, Iterable, Optional
 import re
+from typing import TYPE_CHECKING, Dict, Iterable, Optional
 
 from langchain_core.documents import Document
 from langchain_core.graph_vectorstores.links import get_links
@@ -18,7 +18,8 @@ def _escape_id(id: str) -> str:
     "bidir": "both",
 }
 
-_WORD_RE = re.compile('\s*\S+')
+_WORD_RE = re.compile("\s*\S+")
+
 
 def _split_prefix(s: str, max_chars: int = 50) -> str:
     words = _WORD_RE.finditer(s)
@@ -29,11 +30,12 @@ def _split_prefix(s: str, max_chars: int = 50) -> str:
             break
         split = word.end(0)
 
-    if split ==  len(s):
+    if split == len(s):
         return s
     else:
         return f"{s[0:split]}..."
 
+
 def render_graphviz(
     documents: Iterable[Document],
     engine: Optional[str] = None,
@@ -90,7 +92,10 @@ def render_graphviz(
         escaped_id = _escape_id(id)
         color = node_colors[id] if id in node_colors else node_color
 
-        node_label=f"{graphviz.escape(id)}\n{graphviz.escape(_split_prefix(document.page_content))}"
+        node_label = "\n".join([
+            graphviz.escape(id),
+            graphviz.escape(_split_prefix(document.page_content)),
+        ])
         graph.node(
             escaped_id,
             label=node_label,

libs/community/tests/unit_tests/graph_vectorstores/test_visualize.py

@@ -39,7 +39,8 @@ def test_visualize_simple_graph():
         "\ta -> href_a [dir=back]\n"
         '\tkw_foo [label="kw:foo"]\n'
         "\ta -> kw_foo [dir=both]\n"
-        '\tb [label="b\n<some\n more content>" shape=note tooltip="<some\n more content>"]\n'
+        '\tb [label="b\n<some\n more content>" '
+        'shape=note tooltip="<some\n more content>"]\n'
         '\thref_b [label="href:b"]\n'
         "\tb -> href_b [dir=back]\n"
         "\tb -> href_a [dir=forward]\n"
@@ -55,12 +56,14 @@ def test_visualize_simple_graph():
         "digraph {\n"
         "\trankdir=LR\n"
         "\tnode [style=filled]\n"
-        '\ta [label="a\nsome content" fillcolor=gold shape=note tooltip="some content"]\n'
+        '\ta [label="a\nsome content" fillcolor=gold '
+        'shape=note tooltip="some content"]\n'
         '\thref_a [label="href:a"]\n'
         "\ta -> href_a [dir=back]\n"
         '\tkw_foo [label="kw:foo"]\n'
         "\ta -> kw_foo [dir=both]\n"
-        '\tb [label="b\n<some\n more content>" shape=note tooltip="<some\n more content>"]\n'
+        '\tb [label="b\n<some\n more content>" '
+        'shape=note tooltip="<some\n more content>"]\n'
         '\thref_b [label="href:b"]\n'
         "\tb -> href_b [dir=back]\n"
         "\tb -> href_a [dir=forward]\n"
@@ -70,7 +73,9 @@ def test_visualize_simple_graph():
         "}\n"
     )
 
-    assert render_graphviz([doc1, doc2], node_color = "gold", node_colors={"a": None}).source == (
+    assert render_graphviz(
+        [doc1, doc2], node_color="gold", node_colors={"a": None}
+    ).source == (
         "digraph {\n"
         "\trankdir=LR\n"
         "\tnode [style=filled]\n"
@@ -79,7 +84,8 @@ def test_visualize_simple_graph():
         "\ta -> href_a [dir=back]\n"
         '\tkw_foo [label="kw:foo"]\n'
         "\ta -> kw_foo [dir=both]\n"
-        '\tb [label="b\n<some\n more content>" fillcolor=gold shape=note tooltip="<some\n more content>"]\n'
+        '\tb [label="b\n<some\n more content>" fillcolor=gold '
+        'shape=note tooltip="<some\n more content>"]\n'
         '\thref_b [label="href:b"]\n'
         "\tb -> href_b [dir=back]\n"
         "\tb -> href_a [dir=forward]\n"