Capture image of window in OS X

capture-window.rkt

#lang racket
;;;
;;; How to take a screenshot on OS X
;;;   -- Jens Axel Søgaard
;;;
;;; This gist illustrates how to work with OS X
;;; system calls in order to take a screen shot.
;;;

(require racket/draw
         ffi/unsafe
         ffi/unsafe/objc
         ffi/unsafe/atomic
         ffi/unsafe/nsstring
         mred/private/wx/cocoa/image
         mred/private/wx/cocoa/types)

(import-class NSImage)

; The function CGDisplayCreateImage provided by the Quartz framework,
; will take a screenshot and produce an CGImage.
(define quartz-lib (ffi-lib "/System/Library/Frameworks/Quartz.framework/Versions/Current/Quartz"))

; Each display attached to the computer has a DisplayID.
; We need the id of the display, we want to take a screenshot of.
; In most cases we want a screenshot of the main display.
; An id is represented as an unsigned 32 bit integer.

; typedef uint32_t CGDirectDisplayID;
(define _CGDirectDisplayID _uint32)

(define CGMainDisplayID
  (get-ffi-obj "CGMainDisplayID" quartz-lib  (_fun -> _CGDirectDisplayID)
    (lambda () (error 'quartz-lib "CGMainDisplayID not found"))))

(define (main-display-id)
  (CGMainDisplayID))

; If we need to make screenshots of the other displays,
; we need first to get hold of the ids of the active displays.
(define CGGetActiveDisplayList
  (get-ffi-obj "CGGetActiveDisplayList" quartz-lib 
               (_fun _uint32 (_or-null _pointer) _pointer -> _CGDirectDisplayID)
    (lambda () (error 'quartz-lib "CGGetActiveDisplayList not found"))))

(define (number-of-active-displays)
  ; passing NULL as activeDspys will store
  ; the number of active displays in d
  (define d (malloc _uint32))
  (CGGetActiveDisplayList 0 #f d)
  (ptr-ref d _uint32 0))

(define (get-active-display-ids)
  ; get the number of active displays
  (define n (number-of-active-displays))
  (define d (malloc _uint32))
  ; prepreare am array to hold the display ids
  (define displays (malloc _CGDirectDisplayID n))
  (memset displays 0 (* n (ctype-sizeof _CGDirectDisplayID)))
  (CGGetActiveDisplayList n displays d)
  ; return the display ids a list
  (for/list ([i (in-range n)])
    (ptr-ref displays _CGDirectDisplayID i)))

; The screenshot function returns a pointer to an CGImage:
;     typedef struct CGImage *CGImageRef;
(define _CGImageRef (_cpointer 'CGImageRef))

; The screenshot function:
;     CGImageRef CGDisplayCreateImage( CGDirectDisplayID displayID );
(define CGDisplayCreateImage
  (get-ffi-obj "CGDisplayCreateImage" quartz-lib
               (_fun _CGDirectDisplayID -> _CGImageRef)))

(define (cgimage-screenshot display-id)
  ; take a screenshot of the display with the given id  
  (CGDisplayCreateImage display-id))

(define (cgimage->nsimage cgimage)
  ; convert the CGimage to a NSImage
  (tell (tell NSImage alloc) 
        initWithCGImage: #:type _CGImageRef cgimage 
        size: #:type _NSSize (make-NSSize 0 0)))

(define (screenshot [display-id (main-display-id)])
  ; take a screenshot and convert it into a bitmap%
  (image->bitmap
   (cgimage->nsimage
    (cgimage-screenshot
     (main-display-id)))))

(require slideshow/pict)
(define (scaled-screenshot)
  (scale (bitmap (screenshot)) 0.5))

;;;
;;; Screenshots of windows
;;;

; The function CGWindowListCopyWindowInfo returns
; a CFArray of dictionaries containing information
; on the windows. 

;;; Arrays

; CF (CoreFoundation) has arrays.
(define _CFArrayRef (_cpointer 'CFArrayRef))

; typedef signed long CFIndex;
(define _CFIndex _long)

(define CFArrayGetValueAtIndex
  (get-ffi-obj "CFArrayGetValueAtIndex" quartz-lib
               (_fun _CFArrayRef _CFIndex 
                     -> _pointer)))

(define CFArrayGetCount
  (get-ffi-obj "CFArrayGetCount" quartz-lib
               (_fun _CFArrayRef -> _CFIndex)))

(define (cfarray-length arr)
  (CFArrayGetCount arr))

(define (cfarray-get array idx)
  ; Retrieves a value at a given index.
  (CFArrayGetValueAtIndex array idx))

;;; Window Operations

; typedef uint32_t CGWindowID;
(define _CGWindowID _uint32)

; typedef uint32_t CGWindowListOption;
(define _CGWindowListOption _uint32)

; CGWindowListOption constants:
(define (<< a b) (arithmetic-shift a b))
(define kCGWindowListOptionAll 0)
(define kCGWindowListOptionOnScreenOnly         (<< 1 0))
(define kCGWindowListOptionOnScreenAboveWindow  (<< 1 1))
(define kCGWindowListOptionOnScreenBelowWindow  (<< 1 2))
(define kCGWindowListOptionIncludingWindow      (<< 1 3))
(define kCGWindowListExcludeDesktopElements     (<< 1 4))


(define CGWindowListCopyWindowInfo
  (get-ffi-obj "CGWindowListCopyWindowInfo" quartz-lib
               (_fun _CGWindowListOption _CGWindowID 
                     -> _CFArrayRef)))

;;; Dictionaries

(define _CFDictionaryRef (_cpointer 'CFDictionaryRef))

(define CFDictionaryGetValue
  (get-ffi-obj "CFDictionaryGetValue" quartz-lib
               (_fun _CFDictionaryRef   ; theDict
                     _pointer           ; const void *key
                     -> 
                     _pointer           ; 
                     )))

; Get the number of keys in a dictionary:
;     CFIndex CFDictionaryGetCount ( CFDictionaryRef theDict );
(define CFDictionaryGetCount
  (get-ffi-obj "CFDictionaryGetCount" quartz-lib
               (_fun _CFDictionaryRef   ; theDict
                     -> _CFIndex)))

;void CFDictionaryGetKeysAndValues (
;   CFDictionaryRef theDict,
;   const void * *keys,
;   const void * *values );
(define _CFStringRef (_cpointer 'CFStringRef))
(define _ArrayOfCFStringRef (_cpointer _CFStringRef))
(define CFDictionaryGetKeysAndValues
  (get-ffi-obj "CFDictionaryGetKeysAndValues" quartz-lib
               (_fun _CFDictionaryRef   ; theDict
                     _ArrayOfCFStringRef
                     _ArrayOfCFStringRef
                     -> _CFIndex)))


; CGWindowListCreate
;     Returns the list of window IDs associated with 
;     the specified windows in the current user session.
; CFArrayRef CGWindowListCreate(
;    CGWindowListOption option,
;    CGWindowID relativeToWindow
; );
(define CGWindowListCreate
  (get-ffi-obj "CGWindowListCreate" quartz-lib
               (_fun _CGWindowListOption _CGWindowID 
                     -> _CFArrayRef)))



; a guaranteed invalied WindowId
(define kCGNullWindowID 0)

(define (info-all-windows)
  ; Return a CFArray of CFDictionaryRef types, each of 
  ; which contains information about one of the windows 
  ; in the current user session. If there are no windows 
  ; matching the desired criteria, the function returns 
  ; an empty array. If you call this function from outside 
  ; of a GUI security session or when no window server 
  ; is running, this function returns NULL.
  (CGWindowListCopyWindowInfo kCGWindowListOptionAll kCGNullWindowID))

(define (id-all-windows)
  (CGWindowListCreate kCGWindowListOptionAll kCGNullWindowID))

(define (get-window-id arr idx)
  ; Note: a WindowId is a _uint32,
  ;       but cfarray-get returns
  ;       the value as a _uint64
  ;       (I think)
  (cast (cfarray-get arr idx)
        _pointer _uint64))


; struct CGPoint { CGFloat x; CGFloat y; };
(define-cstruct _CGPoint ([x _CGFloat] [y _CGFloat]))
; struct CGSize { CGFloat width; CGFloat height; };
(define-cstruct _CGSize ([width _CGFloat] [height _CGFloat]))

; struct CGRect { CGPoint origin; CGSize size; };
(define-cstruct _CGRect ([origin _CGPoint] [size _CGSize]))

(define _CGWindowImageOption _uint32)
(define kCGWindowImageDefault             0)
(define kCGWindowImageBoundsIgnoreFraming (<< 1 0))
(define kCGWindowImageShouldBeOpaque      (<< 1 1))
(define kCGWindowImageOnlyShadows         (<< 1 2))
   
(define CGWindowListCreateImage
  (get-ffi-obj "CGWindowListCreateImage" quartz-lib
               (_fun _CGRect              ; screenbounds
                     _CGWindowListOption  ; windowOption
                     _CGWindowID          ; windowId
                     _CGWindowImageOption ; imageOption
                     ->
                     _CGImageRef)))
(define CGRectNull (make-CGRect (make-CGPoint 0 0)
                                (make-CGSize 0 0)))



;;;
;;; Code beyound this point is experimental
;;;

(define windows (CGWindowListCreate (bitwise-ior kCGWindowListOptionOnScreenOnly)
                                    kCGNullWindowID))
(cond
  [(zero? (cfarray-length windows))
   (displayln "No windows matched the criteria!")]
  [else
   (get-window-id windows 0)
   (image->bitmap
    (cgimage->nsimage
     (CGWindowListCreateImage CGRectNull 
                              (bitwise-ior
                               kCGWindowListOptionOnScreenOnly
                               kCGWindowListExcludeDesktopElements
                               kCGWindowListOptionIncludingWindow)
                              (get-window-id windows 0)
                              (bitwise-ior
                               kCGWindowImageDefault
                               ))))])

(define (get-keys-and-values dict)  
  (define n (CFDictionaryGetCount dict))
  (define keys (cast (malloc _CFStringRef n) _pointer _ArrayOfCFStringRef))
  (define vals (cast (malloc _CFStringRef n) _pointer _ArrayOfCFStringRef))
  (memset keys 0 (* n 8))
  (memset vals 0 (* n 8))
  (CFDictionaryGetKeysAndValues dict keys vals)
  ; There is something wrong with the pointer to pointer thing here...
  #;(for/list ([i (in-range n)])
      (list (cast (array-ref (ptr-ref keys (_array _CFStringRef n)) i) _pointer _NSString)
            (cast (array-ref (ptr-ref vals (_array _CFStringRef n)) i) _pointer _NSString))))

(get-keys-and-values
 (cast (cfarray-get 
        (CGWindowListCopyWindowInfo
         (bitwise-ior kCGWindowListOptionOnScreenOnly
                      kCGWindowListExcludeDesktopElements
                      kCGWindowListOptionIncludingWindow)
         (get-window-id windows 0))
        0)
       _pointer _CFDictionaryRef))