source: qesdi/geoplot/trunk/lib/geoplot/plot_writer.py @ 5710

Subversion URL: http://proj.badc.rl.ac.uk/svn/ndg/qesdi/geoplot/trunk/lib/geoplot/plot_writer.py@5710
Revision 5710, 9.0 KB checked in by pnorton, 10 years ago (diff)

Cleaned up a the import statements for the geoplot modules, also fixed a couple of the tests.

Line 
1"""
2plot_writer.py
3==============
4Contains a class responsible for writing a completed plot to an image file.
5"""
6import os
7
8import StringIO
9import logging
10
11from matplotlib.backends.backend_agg import FigureCanvasAgg
12from matplotlib.backends.backend_ps import FigureCanvasPS
13from matplotlib.backends.backend_pdf import FigureCanvasPdf
14from matplotlib.backends.backend_svg import FigureCanvasSVG
15
16
17import image_utils.convert as convert
18#import image_utils.merge_bitmap_logo as merge_bitmap_logo
19import image_utils.merge_eps_logo as merge_eps_logo
20#import image_utils.cmyk_conversion.cmykify_ps as cmykify_ps
21
22#import resources
23from geoplot.image_import import Image
24
25import geoplot.config as geoplot_config
26config = geoplot_config.getConfig()
27
28log = logging.getLogger(__name__)
29
30PLOT_SECTION = 'PlotProperties'
31POSITION_LOGO_IN_CENTER = True
32
33class PlotWriter(object):
34    """
35    A class that takes a plotarea after the plot has been done and writes it
36    to an image file on the local disk.
37    """
38
39    def __init__(self, format, outputCMYK):
40        """
41        Constructs a PlotWriter object
42
43        @keyword format:the format of the output file
44        @type format: a string , one of ('jpg','pdf','ps','png')
45        """
46
47        self.format = format
48        self.outputCMYK = outputCMYK
49
50    def writeFile(self, figure, filename, drawLogo, logoAx=None):
51        """
52        Writes the current plot to the specified filename
53
54        Because mapplotlib cant write to jpg directly first a temporary
55        png file is written then this file is converted into a jpg.
56
57        @param filename: the name of the file (including path) to write the plot to
58        @type filename: string
59        """
60
61        self.checkFileExtension(filename)
62
63        log.info("Writing File :" + filename)
64
65        if self.format in ['png', 'jpg']:
66            initialFormat = 'png'
67        elif self.format in ['ps', 'pdf']:
68            initialFormat = 'ps'
69        else:
70            raise Exception("Unknown format %s" % (self.format,))
71           
72        if initialFormat != self.format:
73            initialFileName = self._getTempFilename(filename, initialFormat)
74        else:
75            initialFileName = filename
76
77
78
79        if drawLogo and initialFormat == 'png':
80                self._addLogo(figure, logoAx)
81
82        self._printFigure(initialFileName, figure)
83       
84        #convert to cmyk (if postscript)
85        if self.outputCMYK and initialFormat == 'ps':
86            pass
87            #self._convertEPSToCMYK(initialFileName)
88       
89        # add the logo
90        if drawLogo and initialFormat == 'ps':
91                if self.outputCMYK:
92                    self._addLogoToPlot(initialFileName, version='CMYK')
93                else:
94                    self._addLogoToPlot(initialFileName, version='RGB')               
95
96       
97        if self.format == 'pdf':
98            pass
99            #convert the .ps to a .pdf
100            convert.epsToPdf(initialFileName, outputFile=filename)
101            os.remove(initialFileName) # remove the old ps file
102                   
103        elif self.format == 'jpg':
104            convert.pngToJpg(initialFileName, filename)
105            os.remove(initialFileName)           
106 
107
108
109    def checkFileExtension(self, filename):
110        """
111        Checks the extenstion of the filename matches the self.Format value
112
113        @param filename:the name of the file to check the format of
114        @type filename:string
115        """
116
117        froot, ext = os.path.splitext(filename)
118        if ext[1:] != self.format :
119            raise Exception("File extension [%s] must match format [%s]!" \
120                             % (ext[1:], self.format))
121
122
123    def _addLogo(self, figure, logoAxes):
124        logo = Image.open(config[PLOT_SECTION]['logoPath'])
125        logo = logo.transpose(Image.FLIP_TOP_BOTTOM)
126        logoAxes.imshow(logo)
127
128    def _addBitmapLogo(self, figure, filename):
129        """
130        Adds a bitmap logo to a .png file. The logo is added in the logoRegion, read from config.ini
131
132        if POSITION_LOGO_IN_CENTER is True then the logo will be centered in the region, otherwise it
133        will be drawn in the top left corner. The figure is needed to caculate the absolute size of the
134        resulting image.
135
136        @param figure: the figure representing the plot
137        @type  figure: a matplotlib.figure instance
138        @param filename: the name of the file to join the logo to
139        @type  filename: string
140        """
141        logoRegion = config['PlotRegions']['logoRegion']
142
143        size = figure.get_size_inches() # get the size in inches
144        size = (size[0] * figure.get_dpi(), size[1] * figure.get_dpi()) #inches * dpi = pixels
145        pxLeft = size[0] * logoRegion[0]
146        # need to transform as logoRegion[1] is distance from bottom not top
147        pxDown = size[1] * (1 - logoRegion[1] - logoRegion[3])
148        maxWidth = size[0] * logoRegion[2]
149        maxHeight = size[1] * logoRegion[3]
150
151        merge_bitmap_logo.stampBitmapLogo(filename, pxLeft, pxDown, maxWidth, maxHeight,
152                                          centerImage=POSITION_LOGO_IN_CENTER,
153                                          printMessages=False)
154
155    def _getTempFilename(self, filename, extension):
156        """
157        Writes the output from a plotArea object to a temporary image file.
158
159        @param  plotArea: the plot area object used to generate the image
160        @type   plotArea: a PlotArea object
161        @param  filename: the filename for the final output file, the name of the
162            temporary file will be based on this name.
163        @type   filename: string
164        @param extension: The image format of the tempoarary file (note this may not
165            be the same as the final output image)
166        @type  extension: string, one of ('jpg', 'pdf,' 'ps', 'png')
167        @return         : the filename of the temporary file that was written
168        @rtype          : string
169        """
170
171        froot, ext = os.path.splitext(filename)
172        tempFileName = froot + "_tmp" + "." + extension
173#        log.debug("temp file name:" + tempFileName)
174        return tempFileName
175
176
177    def _addLogoToPlot(self, filename, version='RGB'):
178
179        if version == 'RGB':
180            logoPath = config[PLOT_SECTION]['EPSLogoPathRGB']
181#            logoFilename = pkg_resources.resource_filename('geoplot',
182#                            config[PLOT_SECTION]['EPSLogoPathRGB'])
183        elif version == 'CMYK':
184            logoPath = config[PLOT_SECTION]['EPSLogoPathCMYK']
185#            logoFilename = pkg_resources.resource_filename('geoplot',
186#                            config[PLOT_SECTION]['EPSLogoPathCMYK'])
187        else:
188            raise Exception("Unknown logo version :" + str(version))
189
190#        logoFilename = resources.getPath(logoPath)
191        logoRegion = config['PlotRegions']['logoRegion']
192
193        log.debug("logoPath = %s" % (logoPath,))
194
195        froot, ext = os.path.splitext(filename)
196
197        outputFile = froot + '_with_logo' + ext
198
199        #join the temporary file with the logo postscript file to produce the ouput file
200        merge_eps_logo.addLogoToFile(filename, outputFile, logoPath, logoRegion)
201
202        log.debug("Removing the old file %s" % (filename, ))
203       
204        #remove the old file
205        os.remove(filename)
206
207        log.debug("Renaming %s to %s" % (outputFile, filename,))
208        #rename the new file to the old one
209        os.rename(outputFile, filename)
210
211    def _convertEPSToCMYK(self, filename):
212
213        froot, ext = os.path.splitext(filename)
214        outputFile = froot + '_cmyk' + ext
215
216        cmykify_ps.cmykifyPostscript(filename, outputFile)
217
218        #remove the old file
219        os.remove(filename)
220
221        #rename the new file to the old one
222        os.rename(outputFile, filename)
223
224    def printFigureToBuffer(self, figure):
225        """
226        Plots the figure directly to StringIO buffer
227        """
228       
229        buffer = StringIO.StringIO()
230        self._printFigure(buffer, figure)
231        buffer.seek(0)
232       
233        return buffer     
234
235    def _printFigure(self, filename, figure):
236        """
237        Instantiate a canvas class and call its print_figure method.
238
239        This is a convenience function that hides a little object assembly.
240
241        @param filename: The output filename
242            (first argument passed to canvas.print_figure())
243        @param CanvasClass: A backend-specific subclass of FigureCanvasBase
244        @param kwargs: Arguments passed to canvas.print_figure()
245
246        """
247
248        if self.format == 'ps':
249            CanvasClass=FigureCanvasPS
250        elif self.format == 'png':
251            CanvasClass=FigureCanvasAgg
252        elif self.format == 'pdf':
253            CanvasClass=FigureCanvasPdf
254        elif self.format == 'svg':
255            CanvasClass = FigureCanvasSVG
256        else:
257            raise Exception("Unknown format to write to , " + self.format)
258
259        canvas = CanvasClass(figure)
260
261        return canvas.print_figure(filename, 
262                                   dpi=figure.get_dpi(),
263                                   facecolor=figure.get_facecolor(), 
264                                   edgecolor=figure.get_edgecolor())
265   
266   
267   
Note: See TracBrowser for help on using the repository browser.