Source code for plantuml

"""
"""
import zlib
import httplib2
import os
import sys

__version__ = [0,1,0]
__version_string__ = '.'.join(str(x) for x in __version__)

__author__ = 'Doug Napoleone'
__email__ = 'doug.napoleone+plantuml@gmail.com'

#: Default plantuml service url
SERVER_URL = 'http://www.plantuml.com/plantuml/img/'

[docs]class PlantUMLError(Exception): """Error in processing. """ pass
[docs]class PlantUMLConnectionError(PlantUMLError): """Error connecting or talking to PlantUML Server. """ pass
[docs]class PlantUMLHTTPError(PlantUMLConnectionError): """Request to PlantUML server returned HTTP Error. """ def __init__(self, response, content, *args, **kwdargs): super(PlantUMLConnectionError, self).__init__(*args, **kwdargs) self.response = response self.content = content if not self.message: self.message = "%d: %s" % ( self.response.status, self.response.reason)
[docs]def deflate_and_encode(plantuml_text): """zlib compress the plantuml text and encode it for the plantuml server. """ zlibbed_str = zlib.compress(plantuml_text) compressed_string = zlibbed_str[2:-4] return encode(compressed_string)
[docs]def encode(data): """encode the plantuml data which may be compresses in the proper encoding for the plantuml server """ res = "" for i in xrange(0,len(data), 3): if (i+2==len(data)): res += _encode3bytes(ord(data[i]), ord(data[i+1]), 0) elif (i+1==len(data)): res += _encode3bytes(ord(data[i]), 0, 0) else: res += _encode3bytes(ord(data[i]), ord(data[i+1]), ord(data[i+2])) return res
def _encode3bytes(b1, b2, b3): c1 = b1 >> 2; c2 = ((b1 & 0x3) << 4) | (b2 >> 4); c3 = ((b2 & 0xF) << 2) | (b3 >> 6); c4 = b3 & 0x3F; res = ""; res += _encode6bit(c1 & 0x3F); res += _encode6bit(c2 & 0x3F); res += _encode6bit(c3 & 0x3F); res += _encode6bit(c4 & 0x3F); return res; def _encode6bit(b): if b < 10: return chr(48 + b) b -= 10 if b < 26: return chr(65 + b) b -= 26 if b < 26: return chr(97 + b); b -= 26 if b == 0: return '-' if b == 1: return '_' return '?'
[docs]class PlantUML(object): """ """ def __init__(self, url=SERVER_URL, basic_auth={}, form_auth={}, http_opts={}, request_opts={}): self.url = url self.request_opts = request_opts self.auth_type = 'basic_auth' if basic_auth else ( 'form_auth' if form_auth else None) self.auth = basic_auth if basic_auth else ( form_auth if form_auth else None) self.http = httplib2.Http(**http_opts) if self.auth_type == 'basic_auth': self.http.add_credentials( self.auth['username'], self.auth['password']) elif self.auth_type == 'form_auth': if 'url' not in self.auth: raise PlantUMLError( "The form_auth option 'url' must be provided and point to " "the login url.") if 'body' not in self.auth: raise PlantUMLError( "The form_auth option 'body' must be provided and include " "a dictionary with the form elements required to log in. " "Example: form_auth={'url': 'http://example.com/login/', " "'body': { 'username': 'me', 'password': 'secret'}") login_url = self.auth['url'] body = self.auth['body'] method = self.auth.get('method', 'POST') headers = self.auth.get( 'headers',{'Content-type':'application/x-www-form-urlencoded'}) try: response, content = http.request( login_url, method, headers=headers, body=urllib.urlencode(body)) except HttpLib2Error, e: raise PlantUMLConnectionError(e) if response.status != 200: raise PlantUMLHTTPError(response, content) self.request_opts['Cookie'] = response['set-cookie']
[docs] def processes(self, plantuml_text): """Processes the plantuml text into the raw PNG image data. """ url = self.url + deflate_and_encode(plantuml_text) try: response, content = self.http.request(url, **self.request_opts) except HttpLib2Error, e: raise PlantUMLConnectionError(e) if response.status != 200: raise PlantUMLHTTPError(response, content) return content
[docs] def processes_file(self, filename, outfile=None, errorfile=None): """Take a filename of a file containing plantuml text and processes it into a .png image. """ if outfile is None: outfile = os.path.splitext(filename)[0] + '.png' if errorfile is None: errorfile = os.path.splitext(filename)[0] + '_error.html' data = open(filename, 'U').read() try: content = self.processes(data) except PlantUMLHTTPError, e: err = open(errorfile, 'w') err.write(e.content) err.close() return False out = open(outfile, 'wb') out.write(content) out.close() return True
if __name__ == '__main__': pl = PlantUML() for filename in sys.argv[1:]: print filename+':', if pl.processes_file(filename): print 'success.' else: print 'failure.'