trivial/format.rkt
2015-12-14 02:36:53 -05:00

102 lines
3.5 KiB
Racket

#lang typed/racket/base
;; Statically-checked format strings
(provide
format:
;; (-> (x:String) Any *N Void)
;; Takes 1 required string argument and N additional arguments,
;; where N is the number of format sequences in the string.
;;
;; If the string is a literal, raises a compile-time arity error if
;; the given number of arguments does not match the format string.
;;
;; If the string is not a literal, arity-checking happens at runtime.
printf:
;; (-> (x:String) Any *N Void)
;; Similar to `format`, but displays the formatted string to `current-output-port`.
)
(require
(for-syntax
typed/racket/base
syntax/parse
racket/sequence))
;; =============================================================================
(define-syntax format:
(syntax-parser
[(f template:str arg* ...)
;; 1. Parse expected types from the template
(let* ([type* (template->type* (syntax-e #'template) #:src #'f)]
[num-expected (length type*)]
[num-given (for/sum ([a (in-syntax #'(arg* ...))]) 1)])
(unless (= num-expected num-given)
(apply raise-arity-error
(syntax-e #'f)
num-expected
(for/list ([a (in-syntax #'(arg* ...))]) (syntax->datum a))))
;; 2. If any types left obligations, use `ann` to typecheck the args
(let ([arg+*
(for/list ([a (in-syntax #'(arg* ...))]
[t (in-list type*)])
(if t (quasisyntax/loc #'f (ann #,a #,t)) a))])
(quasisyntax/loc #'f
(format template #,@arg+*))))]
[f:id
(syntax/loc #'f format)]
[(f tmp arg* ...)
(syntax/loc #'f (format tmp arg* ...))]))
;; Short for `(displayln (format: ...))`
(define-syntax printf:
(syntax-parser
[f:id
(syntax/loc #'f printf)]
[(f arg* ...)
(syntax/loc #'f (display (format: arg* ...)))]))
;; -----------------------------------------------------------------------------
;; Count the number of format escapes in a string.
;; Returns a list of optional types (to be spliced into the source code).
;; Example: If result is '(#f Integer), then
;; - The format string expects 2 arguments
;; - First argument has no type constraints, second must be an Integer
;; (: count-format-escapes (->* [String] [#:src (U #f Syntax)] (Listof (U #f Syntax))))
(define-for-syntax (template->type* str #:src [stx #f])
(define last-index (- (string-length str) 1))
(let loop ([i 0] [acc '()])
(cond
[(>= i last-index)
(reverse acc)]
[(eq? #\~ (string-ref str i))
;; From fprintf docs @ http://docs.racket-lang.org/reference/Writing.html
(case (string-ref str (+ i 1))
[(#\% #\n #\~ #\space #\tab #\newline)
;; Need 0 arguments
(loop (+ i 2) acc)]
[(#\a #\A #\s #\S #\v #\V #\e #\E)
;; Need 1 argument, can be anything
(loop (+ i 2) (cons #f acc))]
[(#\.)
;; Need at most 1, can be anything
(if (and (< (+ 1 i) last-index)
(memq (string-ref str (+ i 2)) '(#\a #\A #\s #\S #\v #\V)))
(loop (+ i 3) (cons #f acc))
(loop (+ i 3) acc))]
[(#\c #\C)
;; Need 1 `char?`
(loop (+ i 2) (cons (syntax/loc stx Char) acc))]
[(#\b #\B #\o #\O #\x #\X)
;; Need 1 `exact?`
(loop (+ i 2) (cons (syntax/loc stx Exact-Number) acc))]
[else
;; Invalid format sequence
(raise-user-error "format: unrecognized pattern string '~~~c'"
(string-ref str (+ i 1)))])]
[else
(loop (+ i 1) acc)])))