dot module refactor

dot.rkt

@@ -1,6 +1,7 @@
 #lang racket/base
 
 (provide
+ dot-graph
  dot-digraph
  dot-vertex
  dot-edge)
@@ -11,26 +12,35 @@
 
 ;; -----------------------------------------------------------------------------
 
-(define (dot-digraph f)
-  (write-string "digraph {\n")
-  (f)
+(define current-dot-edge-syntax (make-parameter "--"))
+
+(define ((dot-graph* header edge-stx) func)
+  (write header)
+  (write-string " {\n")
+  (parameterize ([current-dot-edge-syntax edge-stx])
+    (func))
   (write-string "}\n"))
 
+(define dot-graph (dot-graph* 'graph "--"))
+(define dot-digraph (dot-graph* 'digraph "->"))
+
+(define (write-attrs attrs)
+  (for/fold ([beg "["] [end ""] #:result (write-string end))
+            ([attr (in-list attrs)] #:when (cdr attr))
+    (write-string beg)
+    (printf "~a=~s" (car attr) (cdr attr))
+    (values "," "]")))
+
 (define (dot-vertex id #:label [label #f] #:shape [shape #f])
-  (let* ([params '()]
-         [params (if shape (cons (cons 'shape shape) params) params)]
-         [params (if label (cons (cons 'label label) params) params)])
-    (write id)
-    (for/fold ([beg "["] [end ""] #:result (write-string end))
-              ([param (in-list params)])
-      (write-string beg)
-      (printf "~a=~s" (car param) (cdr param))
-      (values "," "]"))
-    (write-string ";\n")))
+  (write id)
+  (write-attrs `([label . ,label]
+                 [shape . ,shape]))
+  (write-string ";\n"))
 
 (define (dot-edge a b #:flip? [flip? #f])
-  (printf "~a -> ~a;\n"
+  (printf "~a ~a ~a;\n"
           (if flip? b a)
+          (current-dot-edge-syntax)
           (if flip? a b)))
 
 (module+ test
@@ -38,20 +48,39 @@
   (check-equal?
    (with-output-to-string
      (λ ()
-       (dot-digraph
+       (dot-graph
         (λ ()
           (dot-vertex 'a)
           (dot-vertex 'b #:label "B")
           (dot-vertex 'c #:shape 'box)
           (dot-vertex 'd #:label "D" #:shape 'box)))))
    (string-append
-    "digraph {\n"
+    "graph {\n"
     "a;\n"
     "b[label=\"B\"];\n"
     "c[shape=box];\n"
     "d[label=\"D\",shape=box];\n"
     "}\n"))
 
+  (check-equal?
+   (with-output-to-string
+     (λ ()
+       (dot-graph
+        (λ ()
+          (dot-vertex 'a)
+          (dot-vertex 'b)
+          (dot-vertex 'c)
+          (dot-edge 'a 'b)
+          (dot-edge 'b 'c)
+          (dot-edge 'a 'c #:flip? #t)))))
+   (string-append
+    "graph {\n"
+    "a;\nb;\nc;\n"
+    "a -- b;\n"
+    "b -- c;\n"
+    "c -- a;\n"
+    "}\n"))
+
   (check-equal?
    (with-output-to-string
      (λ ()