lisp/batchcolor.lisp @ ad3a9d70d78c

(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* "0.0.1")
(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 ------------------------------------------------------------
(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 "Regex must contain exactly 1 register group, e.g. 'x (fo+) y'."))
    (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 (colorize-line scanner line)))

(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 -----------------------------------------------------------
(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 *option-randomize*
  (adopt:make-option 'randomize
    :help "Randomize the choice of color each run."
    :long "randomize"
    :short #\r
    :reduce (constantly t)))

(defparameter *option-no-randomize*
  (adopt:make-option 'no-randomize
    :result-key 'randomize
    :help "Do not randomize the choice of color each run (the default)."
    :long "no-randomize"
    :short #\R
    :reduce (constantly nil)))

(defparameter *option-debug*
  (adopt:make-option 'debug
    :help "Enable the lisp debugger."
    :long "debug"
    :short #\d
    :reduce (constantly t)))

(defparameter *option-no-debug*
  (adopt:make-option 'no-debug
    :result-key 'debug
    :help "Disable the lisp debugger (the default)."
    :long "no-debug"
    :short #\D
    :reduce (constantly nil)))

(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 *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*)))


(defun toplevel ()
  #+sbcl (sb-ext:disable-debugger)
  (handler-case
      (with-user-abort:with-user-abort
        (multiple-value-bind (arguments options) (adopt:parse-options *ui*)
          (cond
            ((gethash 'help options) (adopt:print-help-and-exit *ui*))
            ((gethash 'version options) (write-line *version*) (adopt:exit))
            (t (if (< (length arguments) 1)
                 (adopt:print-help-and-exit *ui*)
                 (destructuring-bind (pattern . files) arguments
                   (let ((*start* (if (gethash 'randomize options)
                                    (random 256 (make-random-state t))
                                    0)))
                     (run pattern files))))))))
    (with-user-abort:user-abort () (adopt:exit 130))
    (error (c) (adopt:print-error-and-exit c))))