source: trunk/1001Screenshot/screenshot.py@ 604

from CoreGraphics import *
import sys
import os
import subprocess
import tempfile
from Carbon.File import FSPathMakeRef
from Carbon.Folder import FSFindFolder
from Carbon.Folders import kUserDomain, kDesktopFolderType
from Carbon.CF import CFURLCreateFromFileSystemRepresentation
from LaunchServices.Launch import LSSetExtensionHiddenForRef

class HandledError(Exception):
    def __init__(self, args, message, retval):
        # make the arguments look distinct
        for index, arg in enumerate(args):
            if '\\' in arg:
                args[index] = args[index].replace('\\', '\\\\')
            if ' ' in arg:
                args[index] = '"%s"' % args[index].replace('"', '\"')
        description = 'An error occurred while attempting to execute the command %s' % ' '.join(args)
        if message:
            description += ':\n\n%s' % message
        else:
            description += '.\n'
        if retval:
            description += '\nThe process exited with status %d.' % retval
        Exception.__init__(self, description)

def check(args, output, retval):
    if output:
        for index, line in enumerate(output):
            if line.lower().startswith('usage:'):
                output = output[:index]
                break
        raise HandledError(args, ''.join(output), retval)
    elif retval != 0:
        raise HandledError(args, '', retval)

def execute(args):
    process = subprocess.Popen(args, stderr=subprocess.PIPE)
    retval = process.wait()
    check(args, process.stderr.readlines(), retval)

desktopDir = FSFindFolder(kUserDomain, kDesktopFolderType, False).as_pathname()
desktopContents = os.listdir(desktopDir)
index = 1
outFileBase = None
while outFileBase is None:
    outFileBase = 'Screenshot %d' % index
    for filePath in desktopContents:
        if filePath.startswith(outFileBase):
            index += 1
            outFileBase = None
            break
outFileBase = os.path.join(desktopDir, outFileBase)
writtenPaths = []

SCREENCAPTURE_ARGS = ['/usr/sbin/screencapture', '-iW']

outFile = outFileBase + '.png'
args = SCREENCAPTURE_ARGS + ['-tpng', outFile]
screencapture = subprocess.Popen(args, stderr=subprocess.PIPE)
retval = screencapture.wait()
output = screencapture.stderr.readlines()

if retval == 1 and output and 'illegal option -- t' in output[0]:
    # Panther: only PDF output; Finder doesn't notice new files
    def dumpPNG(pdf, pageNumber):
        if pageNumber is None:
            outFile = outFileBase + ".png"
            pageNumber = 1
        else:
            outFile = outFileBase + "-%03d.png" % pageNumber

        w = pdf.getTrimBox(pageNumber).getWidth()
        h = pdf.getTrimBox(pageNumber).getHeight()

        ctx = CGBitmapContextCreateWithColor(int(w), int(h), CGColorSpaceCreateDeviceRGB(), (0,0,0,0))
        ctx.drawPDFDocument(pdf.getTrimBox(pageNumber), pdf, pageNumber)
        ctx.writeToFile(outFile, kCGImageFormatPNG)
        return outFile

    tempDir = tempfile.mkdtemp()
    pdfFile = os.path.join(tempDir, 'screenshot.pdf')

    execute(SCREENCAPTURE_ARGS + [pdfFile])
    if not os.path.exists(pdfFile):
        # something could have gone wrong, or user cancelled, we don't know
        sys.exit(0)

    pdf = CGPDFDocumentCreateWithProvider(CGDataProviderCreateWithFilename(pdfFile))
    numberOfPages = pdf.getNumberOfPages()

    writtenPaths = []
    if numberOfPages == 1:
        writtenPaths = [dumpPNG(pdf, None)]
    else:
        for pageNumber in xrange(1, numberOfPages+1):
            writtenPaths.append(dumpPNG(pdf, pageNumber))

    os.remove(pdfFile)
    os.rmdir(tempDir)

    for path in writtenPaths:
        LSSetExtensionHiddenForRef(FSPathMakeRef(path)[0], True)

    # Carbon.File module doesn't implement FNNotify
    execute(['/usr/bin/osascript', '-e', 'tell app "Finder" to update desktop'])
else:
    check(args, output, retval)
    if not os.path.exists(outFile):
        # something could have gone wrong, or user cancelled, we don't know
        sys.exit(0)
    writtenPaths = [outFile]

# LaunchServices module doesn't implement LSOpenFrom*Spec
execute([os.path.join(os.getenv('RESOURCEPATH'), 'launch'), '-i', 'tv.kungfoo.1001'] + writtenPaths)