1# -*- coding: utf-8 -*-2# Copyright 2011 Takeshi KOMIYA3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.1516from __future__ import division1718import io19import os20import re21from tempfile import NamedTemporaryFile2223from PIL import Image2425from blockdiag.utils import urlutil26from blockdiag.utils.logging import warning2728urlopen_cache = {}293031def urlopen(url, *args, **kwargs):32 """ auto caching urlopen() (using tempfile) """33 from urllib.request import urlopen as orig_urlopen3435 if url not in urlopen_cache:36 with NamedTemporaryFile(delete=False) as tmpfile:37 tmpfile.write(orig_urlopen(url, *args, **kwargs).read())38 tmpfile.flush()39 urlopen_cache[url] = tmpfile.name4041 return io.open(urlopen_cache[url], 'rb')424344def get_image_size(image):45 if isinstance(image, Image.Image):46 return image.size47 else:48 stream = None49 try:50 stream = open(image)51 return stream.size52 finally:53 if stream and hasattr(stream, 'close'):54 stream.close()555657def calc_image_size(size, bounded):58 if bounded[0] < size[0] or bounded[1] < size[1]:59 if (size[0] * 1.0 // bounded[0]) < (size[1] * 1.0 // bounded[1]):60 size = (size[0] * bounded[1] // size[1], bounded[1])61 else:62 size = (bounded[0], size[1] * bounded[0] // size[0])6364 return size656667def color_to_rgb(color):68 import webcolors69 if color == 'none' or isinstance(color, (list, tuple)):70 rgb = color71 elif re.match('#', color):72 rgb = webcolors.hex_to_rgb(color)73 else:74 rgb = webcolors.name_to_rgb(color)7576 return rgb777879def wand_open(url, stream):80 try:81 import wand.image82 except Exception:83 warning("unknown image type: %s", url)84 raise IOError8586 try:87 png_image = io.BytesIO()88 with wand.image.Image(file=stream) as img:89 img.format = 'PNG'90 img.save(file=png_image)91 png_image.seek(0)92 return png_image93 except Exception as exc:94 warning("Fail to convert %s to PNG: %r", url, exc)95 raise IOError969798def pillow_open(url, stream):99 try:100 if isinstance(url, Image.Image):101 return url102 else:103 return Image.open(stream)104 except IOError:105 stream.seek(0)106 png_stream = wand_open(url, stream)107108 return Image.open(png_stream)109110111def open(url, mode='Pillow'):112 if hasattr(url, 'read') or isinstance(url, Image.Image):113 stream = url114 elif not urlutil.isurl(url):115 stream = io.open(url, 'rb')116 else:117 try:118 # wrap BytesIO for rewind stream119 stream = io.BytesIO(urlopen(url).read())120 except Exception:121 warning("Could not retrieve: %s", url)122 raise IOError123124 image = pillow_open(url, stream)125 if mode.lower() == 'pillow':126 # stream will be closed by GC127 return image128 else: # mode == 'png'129 try:130 png_image = io.BytesIO()131 image.save(png_image, 'PNG')132 if hasattr(stream, 'close'): # close() is implemented on Pillow133 stream.close()134 except Exception:135 warning("Could not convert image: %s", url)136 raise IOError137138 png_image.seek(0)139 return png_image140141142def cleanup():143 for url in list(urlopen_cache.keys()):144 path = urlopen_cache.pop(url)145 try:146 os.remove(path)147 except Exception:148 pass149150151def setup(app):152 app.register_cleanup_handler(cleanup)