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

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

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

File size: 4.1 KB
RevLine 
[171]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
[174]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)
[171]28
[174]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)
[171]38
[174]39def execute(args):
40    process = subprocess.Popen(args, stderr=subprocess.PIPE)
41    retval = process.wait()
42    check(args, process.stderr.readlines(), retval)
[171]43
44desktopDir = FSFindFolder(kUserDomain, kDesktopFolderType, False).as_pathname()
45desktopContents = os.listdir(desktopDir)
46index = 1
[174]47outFileBase = None
48while outFileBase is None:
49    outFileBase = 'Screenshot %d' % index
[171]50    for filePath in desktopContents:
[174]51        if filePath.startswith(outFileBase):
[171]52            index += 1
[174]53            outFileBase = None
[171]54            break
[174]55outFileBase = os.path.join(desktopDir, outFileBase)
[171]56writtenPaths = []
57
[174]58SCREENCAPTURE_ARGS = ['/usr/sbin/screencapture', '-iW']
[171]59
[174]60outFile = outFileBase + '.png'
61args = SCREENCAPTURE_ARGS + ['-tpng', outFile]
62screencapture = subprocess.Popen(args, stderr=subprocess.PIPE)
63retval = screencapture.wait()
64output = screencapture.stderr.readlines()
[171]65
[174]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
[171]74
[174]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])
[178]87    if not os.path.exists(pdfFile):
88        # something could have gone wrong, or user cancelled, we don't know
89        sys.exit(0)
[174]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)
[175]111    if not os.path.exists(outFile):
112        # something could have gone wrong, or user cancelled, we don't know
113        sys.exit(0)
[174]114    writtenPaths = [outFile]
115
[171]116# LaunchServices module doesn't implement LSOpenFrom*Spec
[174]117execute([os.path.join(os.getenv('RESOURCEPATH'), 'launch'), '-i', 'tv.kungfoo.1001'] + writtenPaths)
Note: See TracBrowser for help on using the repository browser.