From 4bc6b5699ab23f88a9e95d5a280db038304f218c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Wed, 15 Jun 2022 13:31:40 +0545
Subject: [PATCH] quote node names

---
 dvc/commands/dag.py                | 20 +++++++++++--
 dvc/repo/graph.py                  |  2 +-
 setup.cfg                          |  2 +-
 tests/unit/command/test_dag.py     | 46 ++++++++++++++++++++++++++----
 tests/unit/command/test_imp_url.py |  2 +-
 5 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/dvc/commands/dag.py b/dvc/commands/dag.py
index 502f619e3..20046badb 100644
--- a/dvc/commands/dag.py
+++ b/dvc/commands/dag.py
@@ -1,11 +1,15 @@
 import argparse
+from typing import TYPE_CHECKING
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.ui import ui
 
+if TYPE_CHECKING:
+    from networkx import DiGraph
 
-def _show_ascii(G):
+
+def _show_ascii(G: "DiGraph"):
     from dvc.dagascii import draw
     from dvc.repo.graph import get_pipelines
 
@@ -18,12 +22,24 @@ def _show_ascii(G):
     return "\n".join(ret)
 
 
-def _show_dot(G):
+def _quote_label(node):
+    label = str(node)
+    # Node names should not contain ":" unless they are quoted with "".
+    # See: https://github.com/pydot/pydot/issues/258.
+    if label[0] != '"' and label[-1] != '"':
+        return f'"{label}"'
+    return label
+
+
+def _show_dot(G: "DiGraph"):
     import io
 
+    import networkx as nx
     from networkx.drawing.nx_pydot import write_dot
 
     dot_file = io.StringIO()
+
+    nx.relabel_nodes(G, _quote_label, copy=False)
     write_dot(G.reverse(), dot_file)
     return dot_file.getvalue()
 
diff --git a/dvc/repo/graph.py b/dvc/repo/graph.py
index 9539d7ce5..1f3dc9e81 100644
--- a/dvc/repo/graph.py
+++ b/dvc/repo/graph.py
@@ -33,7 +33,7 @@ def get_pipeline(pipelines, node):
     return found[0]
 
 
-def get_pipelines(G):
+def get_pipelines(G: "DiGraph"):
     import networkx as nx
 
     return [G.subgraph(c).copy() for c in nx.weakly_connected_components(G)]
diff --git a/setup.cfg b/setup.cfg
index b8ce6d667..0f4c220b9 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -51,7 +51,7 @@ install_requires =
     packaging>=19.0
     zc.lockfile>=1.2.1
     flufl.lock>=5
-    networkx>=2.5,<2.8.5
+    networkx>=2.5
     psutil>=5.8.0
     pydot>=1.2.4
     importlib-resources>=5.2.2; python_version < '3.9'
diff --git a/tests/unit/command/test_dag.py b/tests/unit/command/test_dag.py
index 3d115005c..dde967160 100644
--- a/tests/unit/command/test_dag.py
+++ b/tests/unit/command/test_dag.py
@@ -139,18 +139,18 @@ def test_show_dot(repo):
     # output doesn't depend upon order of lines. Use sorted values
     # https://github.com/iterative/dvc/pull/7725
     expected = [
+        "\"stage: '1'\";",
+        "\"stage: '2'\";",
         "\"stage: '3'\" -> \"stage: '4'\";",
+        "\"stage: '3'\";",
+        "\"stage: '4'\";",
         "\"stage: 'a.dvc'\" -> \"stage: '1'\";",
         "\"stage: 'a.dvc'\" -> \"stage: '3'\";",
         "\"stage: 'a.dvc'\" -> \"stage: '4'\";",
+        "\"stage: 'a.dvc'\";",
         "\"stage: 'b.dvc'\" -> \"stage: '2'\";",
         "\"stage: 'b.dvc'\" -> \"stage: '3'\";",
-        "stage;",
-        "stage;",
-        "stage;",
-        "stage;",
-        "stage;",
-        "stage;",
+        "\"stage: 'b.dvc'\";",
         "strict digraph  {",
         "}",
     ]
@@ -160,6 +160,40 @@ def test_show_dot(repo):
     assert actual == expected
 
 
+def test_show_dot_properly_escapes():
+    graph = nx.DiGraph(
+        [
+            ("evaluate", "train🚄"),  # emoji
+            ("evaluate", "featurize"),
+            ("featurize", "prepare:1"),  # colon
+            ("prepare:1", "data/raw/1.dvc"),  # posix path
+            ("prepare:1", "data\\raw\\2.dvc"),  # windows path
+            ("prepare", "4"),  # just a number
+        ]
+    )
+
+    expected = {
+        "strict digraph  {",
+        '"data\\raw\\2.dvc";',
+        '"prepare";',
+        '"4";',
+        '"data/raw/1.dvc";',
+        '"train🚄";',
+        '"evaluate";',
+        '"prepare:1";',
+        '"featurize";',
+        '"data\\raw\\2.dvc" -> "prepare:1";',
+        '"4" -> "prepare";',
+        '"data/raw/1.dvc" -> "prepare:1";',
+        '"train🚄" -> "evaluate";',
+        '"prepare:1" -> "featurize";',
+        '"featurize" -> "evaluate";',
+        "}",
+    }
+    actual = {line.rstrip() for line in _show_dot(graph).splitlines()}
+    assert actual == expected
+
+
 def test_show_mermaid(repo):
     assert [
         line.rstrip() for line in _show_mermaid(repo.index.graph).splitlines()
diff --git a/tests/unit/command/test_imp_url.py b/tests/unit/command/test_imp_url.py
index cfabc388b..3204809c3 100644
--- a/tests/unit/command/test_imp_url.py
+++ b/tests/unit/command/test_imp_url.py
@@ -117,7 +117,7 @@ def test_import_url_to_remote(mocker):
     )
 
 
-def test_import_url_to_remote_invalid_combination(mocker, caplog):
+def test_import_url_to_remote_invalid_combination(dvc, mocker, caplog):
     cli_args = parse_args(
         [
             "import-url",