From 3572318c6d2d77cdc29870f81639aa18a6de84e4 Mon Sep 17 00:00:00 2001
From: Matthew Butterick <mb@mbtype.com>
Date: Tue, 31 Dec 2019 16:32:47 -0800
Subject: [PATCH] make input arg of `attr-ref` more lenient

---
 txexpr/base.rkt                 | 8 +++++---
 txexpr/scribblings/txexpr.scrbl | 9 +++++++--
 txexpr/test/tests.rkt           | 4 ++++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/txexpr/base.rkt b/txexpr/base.rkt
index 80ed2a8..dc1248f 100644
--- a/txexpr/base.rkt
+++ b/txexpr/base.rkt
@@ -235,9 +235,11 @@
   (attr-set tx key (string-join (append starting-values (list value)) " ")))      
 
 (define no-failure-result (gensym)) ; failure-result might be #false
-(define+provide+safe (attr-ref tx key [failure-result no-failure-result])
-  ((txexpr? can-be-txexpr-attr-key?) (any/c) . ->* . any)
-  (match (assq (->txexpr-attr-key key) (get-attrs tx))
+(define+provide+safe (attr-ref attrs-arg key [failure-result no-failure-result])
+  (((or/c txexpr? txexpr-attrs?) can-be-txexpr-attr-key?) (any/c) . ->* . any)
+  (match (assq (->txexpr-attr-key key) (match attrs-arg
+                                         [(? txexpr? tx) (get-attrs tx)]
+                                         [attrs attrs]))
     [(list _ value) value]
     [_ (match failure-result
          [(? procedure?) (failure-result)]
diff --git a/txexpr/scribblings/txexpr.scrbl b/txexpr/scribblings/txexpr.scrbl
index 5a1f62c..d276ba1 100644
--- a/txexpr/scribblings/txexpr.scrbl
+++ b/txexpr/scribblings/txexpr.scrbl
@@ -372,19 +372,24 @@ Return @racket[#t] if @racket[_attrs] and @racket[_other-attrs] contain the same
 
 @defproc[
 (attr-ref
-[tx txexpr?]
+[attrs (or/c txexpr-attrs? txexpr?)]
 [key can-be-txexpr-attr-key?]
 [failure-result any/c (λ _ (raise (make-exn:fail:contract ....)))
 ])
 any]
-Given a @racket[_key], return the corresponding @racket[_value] from the attributes of a @racket[_txexpr]. By default, asking for a nonexistent key produces an error. But if a value or procedure is provided as the @racket[_failure-result], evaluate and return that instead.
+Given a @racket[_key], return the corresponding @racket[_value] from @racket[_attrs]. By default, asking for a nonexistent key produces an error. But if a value or procedure is provided as the @racket[_failure-result], evaluate and return that instead.
 
 @examples[#:eval my-eval
+(define tx '(div ((id "top")(class "red")) "Hello" (p "World")))
 (attr-ref tx 'class)
 (attr-ref tx 'id)
 (attr-ref tx 'nonexistent-key)
 (attr-ref tx 'nonexistent-key "forty-two")
 (attr-ref tx 'nonexistent-key (λ _ (* 6 7)))
+(define attrs '((id "top")(class "red")))
+(attr-ref attrs 'class)
+(attr-ref attrs 'id)
+(attr-ref attrs 'nonexistent-key)
 ]
 
 @deftogether[(
diff --git a/txexpr/test/tests.rkt b/txexpr/test/tests.rkt
index cedbbf8..37128f8 100644
--- a/txexpr/test/tests.rkt
+++ b/txexpr/test/tests.rkt
@@ -150,6 +150,10 @@
  (check-exn exn:fail? (λ _ (attr-ref '(p ((foo "bar"))) 'zam)))
  (check-equal? (attr-ref '(p ((foo "bar"))) 'zam 42) 42)
  (check-equal? (attr-ref '(p ((foo "bar"))) 'zam (λ _ (* 6 7))) 42)
+ (check-equal? (attr-ref '((foo "bar")) 'foo) "bar")
+ (check-exn exn:fail? (λ _ (attr-ref '((foo "bar")) 'zam)))
+ (check-equal? (attr-ref '((foo "bar")) 'zam 42) 42)
+ (check-equal? (attr-ref '((foo "bar")) 'zam (λ _ (* 6 7))) 42)
  (check-txexprs-equal? (attr-set '(p ((foo "bar"))) 'foo "fraw") '(p ((foo "fraw"))))
  (check-txexprs-equal? (attr-set* '(p ((foo "bar"))) 'foo "fraw") '(p ((foo "fraw"))))
  (check-true (let ([result (attr-set* '(p ((foo "bar"))) 'foo "fraw" 'zim 'zam)])