From 2679e6c0724788cba4deae3338ad1acc4e1dd1aa Mon Sep 17 00:00:00 2001
From: jdewar <jeremy.dewar@gmail.com>
Date: Wed, 27 Sep 2023 18:44:45 -0400
Subject: [PATCH] Subgraph outputs DG (#25)

* output of DG.subgraph should be of type DG

* minor spelling fix
---
 lib/dg.ex        | 10 ++++------
 test/dg_test.exs | 23 ++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/lib/dg.ex b/lib/dg.ex
index f9c8e7f..3a65bf9 100644
--- a/lib/dg.ex
+++ b/lib/dg.ex
@@ -205,12 +205,10 @@ defmodule DG do
     :digraph_utils.strong_components(dg)
   end
 
-  def subgraph(%__MODULE__{dg: dg}, vertices) do
-    :digraph_utils.subgraph(dg, vertices)
-  end
-
-  def subgraph(%__MODULE__{dg: dg}, vertices, options) do
-    :digraph_utils.subgraph(dg, vertices, options)
+  def subgraph(%__MODULE__{dg: dg}, vertices, options \\ []) do
+    {digraph_opts, opts} = Keyword.pop(options, :digraph_opts, [])
+    subgraph = :digraph_utils.subgraph(dg, vertices, digraph_opts)
+    %__MODULE__{dg: subgraph, opts: opts}
   end
 
   def topsort(%__MODULE__{dg: dg}) do
diff --git a/test/dg_test.exs b/test/dg_test.exs
index 9e3c0e1..c68e405 100644
--- a/test/dg_test.exs
+++ b/test/dg_test.exs
@@ -112,10 +112,31 @@ defmodule DG.Test do
     end
   end
 
+  describe "subgraph" do
+    test "returns a DG" do
+      dg = DG.new(test: true)
+      ~w(a b c d e) |> Enum.map(&{:vertex, &1}) |> Enum.into(dg)
+
+      assert %DG{} = DG.subgraph(dg, ~w(a b c))
+    end
+
+    test "handles digraph options" do
+      dg = DG.new(test: true)
+      ~w(a b c) |> Enum.map(&{:vertex, &1}) |> Enum.into(dg)
+
+      label = "a -> b"
+      Enum.into([{:edge, "a", "b", label}], dg)
+
+      subgraph = DG.subgraph(dg, ~w(a b), digraph_opts: [keep_labels: true])
+
+      assert {_, "a", "b", ^label} = DG.edge(subgraph, List.first(DG.edges(dg, "a")))
+    end
+  end
+
   describe "sigil" do
     import DG.Sigil
 
-    test "ingetration" do
+    test "integration" do
       dg = ~g"""
       graph LR
       a[aaaaa]