source: trunk/1001Screenshot/screenshot.py

Last change on this file was 178, checked in by Nicholas Riley, 13 years ago

screenshot.py: Properly handle user cancellation in Panther, too.

File size: 4.1 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    if not os.path.exists(pdfFile):
88        # something could have gone wrong, or user cancelled, we don't know
89        sys.exit(0)
90
91    pdf = CGPDFDocumentCreateWithProvider(CGDataProviderCreateWithFilename(pdfFile))
92    numberOfPages = pdf.getNumberOfPages()
93
94    writtenPaths = []
95    if numberOfPages == 1:
96        writtenPaths = [dumpPNG(pdf, None)]
97    else:
98        for pageNumber in xrange(1, numberOfPages+1):
99            writtenPaths.append(dumpPNG(pdf, pageNumber))
100
101    os.remove(pdfFile)
102    os.rmdir(tempDir)
103
104    for path in writtenPaths:
105        LSSetExtensionHiddenForRef(FSPathMakeRef(path)[0], True)
106
107    # Carbon.File module doesn't implement FNNotify
108    execute(['/usr/bin/osascript', '-e', 'tell app "Finder" to update desktop'])
109else:
110    check(args, output, retval)
111    if not os.path.exists(outFile):
112        # something could have gone wrong, or user cancelled, we don't know
113        sys.exit(0)
114    writtenPaths = [outFile]
115
116# LaunchServices module doesn't implement LSOpenFrom*Spec
117execute([os.path.join(os.getenv('RESOURCEPATH'), 'launch'), '-i', 'tv.kungfoo.1001'] + writtenPaths)
Note: See TracBrowser for help on using the repository browser.