lisp/batchcolor.lisp @ a7f0e08891c8

(eval-when (:compile-toplevel :load-toplevel :execute)
  (ql:quickload '(:adopt :cl-ppcre :with-user-abort) :silent t))

(defpackage :batchcolor
  (:use :cl)
  (:export :toplevel :*ui*))

(in-package :batchcolor)

;;;; Configuration ------------------------------------------------------------
(defparameter *version* "1.0.0")
(defparameter *start* 0)

(defun rgb-code (r g b)
  ;; The 256 color mode color values are essentially r/g/b in base 6, but
  ;; shifted 16 higher to account for the intiial 8+8 colors.
  (+ (* r 36)
     (* g 6)
     (* b 1)
     16))

(defparameter *colors* (let ((result (make-array 256 :fill-pointer 0)))
                         (dotimes (r 6)
                           (dotimes (g 6)
                             (dotimes (b 6)
                               (unless (< (+ r g b) 3)
                                 ;; Don't use dark, hard-to-read colors.
                                 (vector-push-extend (rgb-code r g b) result)))))
                         result))


;;;; Functionality ------------------------------------------------------------
(define-condition bad-regex (error) ()
  (:report "Invalid regex."))

(define-condition bad-regex-group-count (bad-regex) ()
  (:report "Invalid regex: must contain exactly 1 register group, e.g. 'x (fo+) y'."))


(defun djb2 (string)
  ;; http://www.cse.yorku.ca/~oz/hash.html
  (reduce (lambda (hash c)
            (mod (+ (* 33 hash) c) (expt 2 64)))
          string
          :initial-value 5381
          :key #'char-code))

(defun find-color (string)
  (aref *colors* (mod (+ (djb2 string) *start*)
                      (length *colors*))))

(defun ansi-color-start (color)
  (format nil "~C[38;5;~Dm" #\Escape color))

(defun ansi-color-end ()
  (format nil "~C[0m" #\Escape))

(defun colorize-line (scanner line &aux (start 0))
  (ppcre:do-scans (ms me rs re scanner line)
    (setf rs (remove nil rs)
          re (remove nil re))
    (when (/= 1 (length rs))
      (error 'bad-regex-group-count))
    (let* ((word-start (aref rs 0))
           (word-end (aref re 0))
           (word (subseq line word-start word-end))
           (color (find-color word)))
      (write-string line *standard-output* :start start :end word-start)
      (format t "~A~A~A" (ansi-color-start color) word (ansi-color-end))
      (setf start word-end)))
  (write-line line *standard-output* :start start)
  (values))


;;;; Run ----------------------------------------------------------------------
(defun run% (scanner stream)
  (loop :for line = (read-line stream nil)
        :while line
        :do (tagbody retry
              (restart-case (colorize-line scanner line)
                (supply-new-regex (v)
                    :test (lambda (c) (typep c 'bad-regex))
                    :report "Supply a new regular expression as a string."
                    :interactive (lambda () (list (read-line *query-io*)))
                  (setf scanner (ppcre:create-scanner v))
                  (go retry))))))

(defun run (pattern paths)
  (if (null paths)
    (setf paths '("-")))
  (let ((scanner (ppcre:create-scanner pattern)))
    (dolist (path paths)
      (if (string= "-" path)
        (run% scanner *standard-input*)
        (with-open-file (stream path :direction :input)
          (run% scanner stream))))))


;;;; User Interface -----------------------------------------------------------
(defmacro defparameters (parameters values-form)
  `(progn
     ,@(loop :for parameter :in parameters
             :collect `(defparameter ,parameter nil))
     (setf (values ,@parameters) ,values-form)
     ',parameters))

(defun make-boolean-options
    (name &key
     (name-no (intern (concatenate 'string (string 'no-) (string name))))
     long
     (long-no (when long (format nil "no-~A" long)))
     short
     (short-no (when short (char-upcase short)))
     (result-key name)
     help
     help-no
     manual
     manual-no
     initial-value)
  (values (adopt:make-option name
            :result-key result-key
            :long long
            :short short
            :help help
            :manual manual
            :initial-value initial-value
            :reduce (constantly t))
          (adopt:make-option name-no
            :result-key result-key
            :long long-no
            :short short-no
            :help help-no
            :manual manual-no
            :reduce (constantly nil))))


(defparameters (*option-randomize* *option-no-randomize*)
  (make-boolean-options 'randomize
    :help "Randomize the choice of color each run."
    :help-no "Do not randomize the choice of color each run (the default)."
    :long "randomize"
    :short #\r))

(defparameters (*option-debug* *option-no-debug*)
  (make-boolean-options 'debug
    :long "debug"
    :short #\d
    :help "Enable the Lisp debugger."
    :help-no "Disable the Lisp debugger (the default)."))

(defparameter *option-help*
  (adopt:make-option 'help
    :help "Display help and exit."
    :long "help"
    :short #\h
    :reduce (constantly t)))

(defparameter *option-version*
  (adopt:make-option 'version
    :help "Display version information and exit."
    :long "version"
    :reduce (constantly t)))


(adopt:define-string *help-text*
  "batchcolor takes a regular expression and matches it against standard ~
   input one line at a time.  Each unique match is highlighted in a random ~
   color.")

(defparameter *examples*
  '(("Colorize IRC nicknames in a chat log:"
     . "cat channel.log | batchcolor '<(\\\\w+)>'")
    ("Colorize UUIDs in a request log:"
     . "tail -f /var/log/foo | batchcolor '([a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12})'")))


(defparameter *ui*
  (adopt:make-interface
    :name "batchcolor"
    :usage "[OPTIONS] REGEX [FILE...]"
    :summary "colorize regex matches in batches"
    :help *help-text*
    :examples *examples*
    :contents (list *option-randomize*
                    *option-no-randomize*
                    *option-debug*
                    *option-no-debug*
                    *option-help*
                    *option-version*)))


(defmacro without-debugger (&body body)
  `(multiple-value-prog1
     (progn
       #+sbcl (sb-ext:disable-debugger)
       (progn ,@body))
     (progn
       #+sbcl (sb-ext:enable-debugger))))

(defmacro exit-on-error (&body body)
  `(without-debugger
     (handler-case (progn ,@body)
       (error (c) (adopt:print-error-and-exit c)))))

(defmacro exit-on-error-unless (expr &body body)
  `(if ,expr
     (progn ,@body)
     (exit-on-error ,@body)))

(defmacro exit-on-ctrl-c (&body body)
  `(handler-case
       (with-user-abort:with-user-abort (progn ,@body))
     (with-user-abort:user-abort () (adopt:exit 130))))


(defun toplevel ()
  (exit-on-ctrl-c
    (multiple-value-bind (arguments options) (adopt:parse-options-or-exit *ui*)
      (exit-on-error-unless (gethash 'debug options)
        (cond
          ((gethash 'help options) (adopt:print-help-and-exit *ui*))
          ((gethash 'version options) (write-line *version*) (adopt:exit))
          ((null arguments) (error "A regular expression is required."))
          (t (destructuring-bind (pattern . files) arguments
               (let ((*start* (if (gethash 'randomize options)
                                (random 256 (make-random-state t))
                                0)))
                 (run pattern files)))))))))