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

Last change on this file since 174 was 174, checked in by Nicholas Riley, 16 years ago

screenshot.py: Much better error handling, including dealing with
screencapture not always reporting errors by exit code
(rdar://problem/4104545). Generate (single) PNGs directly on Tiger -
untested.

File size: 3.8 KB
Line 
1from CoreGraphics import *
2import sys
3import os
4import subprocess
5import tempfile
6from Carbon.File import FSPathMakeRef
7from Carbon.Folder import FSFindFolder
8from Carbon.Folders import kUserDomain, kDesktopFolderType
9from Carbon.CF import CFURLCreateFromFileSystemRepresentation
10from LaunchServices.Launch import LSSetExtensionHiddenForRef
11
12class HandledError(Exception):
13    def __init__(self, args, message, retval):
14        # make the arguments look distinct
15        for index, arg in enumerate(args):
16            if '\\' in arg:
17                args[index] = args[index].replace('\\', '\\\\')
18            if ' ' in arg:
19                args[index] = '"%s"' % args[index].replace('"', '\"')
20        description = 'An error occurred while attempting to execute the command %s' % ' '.join(args)
21        if message:
22            description += ':\n\n%s' % message
23        else:
24            description += '.\n'
25        if retval:
26            description += '\nThe process exited with status %d.' % retval
27        Exception.__init__(self, description)
28
29def check(args, output, retval):
30    if output:
31        for index, line in enumerate(output):
32            if line.lower().startswith('usage:'):
33                output = output[:index]
34                break
35        raise HandledError(args, ''.join(output), retval)
36    elif retval != 0:
37        raise HandledError(args, '', retval)
38
39def execute(args):
40    process = subprocess.Popen(args, stderr=subprocess.PIPE)
41    retval = process.wait()
42    check(args, process.stderr.readlines(), retval)
43
44desktopDir = FSFindFolder(kUserDomain, kDesktopFolderType, False).as_pathname()
45desktopContents = os.listdir(desktopDir)
46index = 1
47outFileBase = None
48while outFileBase is None:
49    outFileBase = 'Screenshot %d' % index
50    for filePath in desktopContents:
51        if filePath.startswith(outFileBase):
52            index += 1
53            outFileBase = None
54            break
55outFileBase = os.path.join(desktopDir, outFileBase)
56writtenPaths = []
57
58SCREENCAPTURE_ARGS = ['/usr/sbin/screencapture', '-iW']
59
60outFile = outFileBase + '.png'
61args = SCREENCAPTURE_ARGS + ['-tpng', outFile]
62screencapture = subprocess.Popen(args, stderr=subprocess.PIPE)
63retval = screencapture.wait()
64output = screencapture.stderr.readlines()
65
66if retval == 1 and output and 'illegal option -- t' in output[0]:
67    # Panther: only PDF output; Finder doesn't notice new files
68    def dumpPNG(pdf, pageNumber):
69        if pageNumber is None:
70            outFile = outFileBase + ".png"
71            pageNumber = 1
72        else:
73            outFile = outFileBase + "-%03d.png" % pageNumber
74
75        w = pdf.getTrimBox(pageNumber).getWidth()
76        h = pdf.getTrimBox(pageNumber).getHeight()
77
78        ctx = CGBitmapContextCreateWithColor(int(w), int(h), CGColorSpaceCreateDeviceRGB(), (0,0,0,0))
79        ctx.drawPDFDocument(pdf.getTrimBox(pageNumber), pdf, pageNumber)
80        ctx.writeToFile(outFile, kCGImageFormatPNG)
81        return outFile
82
83    tempDir = tempfile.mkdtemp()
84    pdfFile = os.path.join(tempDir, 'screenshot.pdf')
85
86    execute(SCREENCAPTURE_ARGS + [pdfFile])
87
88    pdf = CGPDFDocumentCreateWithProvider(CGDataProviderCreateWithFilename(pdfFile))
89    numberOfPages = pdf.getNumberOfPages()
90
91    writtenPaths = []
92    if numberOfPages == 1:
93        writtenPaths = [dumpPNG(pdf, None)]
94    else:
95        for pageNumber in xrange(1, numberOfPages+1):
96            writtenPaths.append(dumpPNG(pdf, pageNumber))
97
98    os.remove(pdfFile)
99    os.rmdir(tempDir)
100
101    for path in writtenPaths:
102        LSSetExtensionHiddenForRef(FSPathMakeRef(path)[0], True)
103
104    # Carbon.File module doesn't implement FNNotify
105    execute(['/usr/bin/osascript', '-e', 'tell app "Finder" to update desktop'])
106else:
107    check(args, output, retval)
108    writtenPaths = [outFile]
109
110# LaunchServices module doesn't implement LSOpenFrom*Spec
111execute([os.path.join(os.getenv('RESOURCEPATH'), 'launch'), '-i', 'tv.kungfoo.1001'] + writtenPaths)
Note: See TracBrowser for help on using the repository browser.