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/'

class PlantUMLError(Exception):
    """Error in processing.
    """
    pass

class PlantUMLConnectionError(PlantUMLError):
    """Error connecting or talking to PlantUML Server.
    """
    pass

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)

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)

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 '?'


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']

    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
    
    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.'