<commit_before>"""
oceanidanalysis.maps
====================

Plotting tools and wrappers to generate common map views of MBARI LRAUV data.

"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

from oceanidanalysis import lib

class Map(Basemap):
    """Subclass of basemap.Basemap with standard colors and boundaries.

    """
    def arangelat(self, dlat=1.):
        try:
            parallels = np.arange(self.llcrnrlat, self.urcrnrlat, dlat)
        except Exception, e:
            # corner attributes dne, probably initialized some other way...
            raise e;
        return parallels

    def arangelon(self, dlon=1.):
        try:
            meridians = np.arange(self.llcrnrlon, self.urcrnrlon, dlon)
        except Exception, e:
            # corner attributes dne, probably initialized some other way...
            raise e;
        return meridians

    def drawparallels(self, circles=None, dlat=1.,
        labels=[1,0,0,0], fontsize=10, **kwargs):
        """Overloaded method with default args.

        """
        kwargs.update(dict(labels=labels, fontsize=fontsize)) #TODO pythonic
        if circles is None: circles = self.arangelat(dlat)
        Basemap.drawparallels(self, circles, **kwargs)

    def drawmeridians(self, circles=None, dlon=1.,
        labels=[0,0,0,1], fontsize=10, **kwargs):
        """Overloaded method with default args.

        """
        kwargs.update(dict(labels=labels, fontsize=fontsize)) #TODO pythonic
        if circles is None: circles = self.arangelon(dlon)
        Basemap.drawmeridians(self, circles, **kwargs)

    def drawgrid(self, lats=None, lons=None, **kwargs):
        """Draw a grid of latitude and longitude lines."""
        self.drawparallels(circles=lats, **kwargs)
        self.drawmeridians(circles=lons, **kwargs)

    def drawdefault(self):
        #TODO accept several kw dicts
        self.fillcontinents(color='coral',lake_color='aqua')
        #self.drawmapboundary(fill_color='aqua')
        self.drawmapboundary()
        self.drawrivers(color='b')

    def draw_bathymetry(self, isobaths):
        """Draw isobaths under water.

        """
        raise NotImplementedError

    def fill_bathymetry(self, cmap):
        """Fill bathymetry under water. (Color plot of sounding depths)

        """
        raise NotImplementedError

    def draw_topography(self, ):
        raise NotImplementedError

    def fill_topography(self, ):
        raise NotImplementedError

    def draw_buoys(self, ):
        """Permanent marker buoys.

        """
        raise NotImplementedError

    def draw_shipping_lanes(self, ):
        """Pretty Self-explanatory.

        """
        raise NotImplementedError

    def draw_track(self, track, **kwargs):
        """Draw a position track on the map (e.g., from Tethys or Daphne).

        """
        raise NotImplementedError
#TODO a track that changes color based on an ancillary variable

    def draw_currents(self, lat, lon, u, v, contourf_magnitude=True, **kwargs):
        """Quiver plot the ocean current.

        """
        if lat.ndim is 2: lat = np.sort(np.unique(lat))
        if lon.ndim is 2: lon = np.sort(np.unique(lon))
        um, vm, xm, ym = self.transform_vector(u, v, lon, lat,
            len(lon), len(lat), returnxy = True, masked = True, order = 1)
        #TODO   pass kwargs to transform_vector?
        #TODO   optional input/output of *m variables above...
        if contourf_magnitude:
            mm = (um**2 + vm**2)**0.5 # 2D current magnitude
            dc = 0.01
#TODO            cmax = max(mm) + dc
            cmax = 0.5 + dc
            mmlevels = np.arange(0, cmax, dc)
            self.contourf(xm, ym, mm, mmlevels, cmap=plt.cm.Blues)
            self.colorbar() #TODO pass kwargs
        q = self.quiver(xm, ym, um, vm, **kwargs)
            #or specify, e.g., width=0.003, scale=400)
        qkey = plt.quiverkey(q, 0.1, -0.2, 1, '1 m/s', labelpos='W')
            #TODO as defaults
        return q, qkey

    def draw_wind(self, wind, altitude, **kwargs):
        """Quiver plot the winds.

        (with or without a mask over the land)
        """
        #TODO redefine args, or unpack wind into u,v,lat,lon
        # transform from spherical to map projection coordinates (rotation
        # and interpolation).
        default = dict(length = 6, linewidth = 0.5,
            barbcolor = 'black', flagcolor = 'cyan', )
        nxv = 25; nyv = 25 #TODO can these be included in defaults?
        default.update(kwargs)
        um, vm, xm, ym = self.transform_vector(u,v,lons,lats,nxv,nyv,returnxy=True)
        self.barbs(xm, ym, um, vm, **default) #TODO what does this return?
        raise NotImplementedError

    def play(self, **kwargs):
        """Generate a movie using inputs with a time dimension.

        """
        raise NotImplementedError

class MontereyBay(Map):
    """A standard map canvas of Monterey Bay.

    """
    def __init__(self, lat_0 = 36.75, lon_0 = -121.0,
        llcrnrlat = 36.5, llcrnrlon = -122.5,
        urcrnrlat = 37, urcrnrlon = -121.75,
        projection = 'tmerc', resolution = 'i', **kwargs):
        """Overloaded method with default args.

        """
        Map.__init__(self, **lib.injectlocals(locals()))
        #TODO some handling if the map is defined in alternate fashion, e.g.:
#m = Basemap(width=920000,height=1100000,
#            resolution='f',projection='tmerc',lon_0=-4.2,lat_0=54.6)

    def drawparallels(self, dlat=.1, **kwargs):
        """Overloaded method with default args.

        """
        Map.drawparallels(self, dlat=dlat, **kwargs)

    def drawmeridians(self, dlon=.1, **kwargs):
        """Overloaded method with default args.

        """
        Map.drawmeridians(self, dlon=dlon, **kwargs)

    def draw_mbari(self, ):
        """Mark location of MBARI on the shore.

        """
        raise NotImplementedError

    def draw_canyon(self, ):
        """Draw boundaries of Monterey Canyon on the map.

        """
        raise NotImplementedError
#TODO other canyons, too?

    def draw_mountains(self, ):
        raise NotImplementedError

    def draw_mars(self,color='red', marker='s', label='MARS', **kwargs):
        """Mark position of MARS cabled observatory.

        """
        plot_kwargs = lib.injectlocals(locals())
        print plot_kwargs
        mars_lat = 36 + 42.7481/60 #TODO use proj or write a utility for oa.lib
        mars_lon =  -(122 + 11.2139/60)
        mars_depth = 891 # meters
        x, y = self(mars_lon, mars_lat) # position in projection
        self.plot(x, y, **plot_kwargs)
        #TODO what figure does it go to by default?

<commit_msg>minor mods<commit_after>"""
oceanidanalysis.maps
====================

Plotting tools and wrappers to generate common map views of MBARI LRAUV data.

"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

from oceanidanalysis import lib

class Map(Basemap):
    """Subclass of basemap.Basemap with standard colors and boundaries.

    """
    def arangelat(self, dlat=1.):
        try:
            parallels = np.arange(self.llcrnrlat, self.urcrnrlat, dlat)
        except Exception, e:
            # corner attributes dne, probably initialized some other way...
            raise e;
        return parallels

    def arangelon(self, dlon=1.):
        try:
            meridians = np.arange(self.llcrnrlon, self.urcrnrlon, dlon)
        except Exception, e:
            # corner attributes dne, probably initialized some other way...
            raise e;
        return meridians

    def drawparallels(self, circles=None, dlat=1.,
        labels=[1,0,0,0], fontsize=10, **kwargs):
        """Overloaded method with default args.

        """
        kwargs.update(dict(labels=labels, fontsize=fontsize)) #TODO pythonic
        if circles is None: circles = self.arangelat(dlat)
        Basemap.drawparallels(self, circles, **kwargs)

    def drawmeridians(self, circles=None, dlon=1.,
        labels=[0,0,0,1], fontsize=10, **kwargs):
        """Overloaded method with default args.

        """
        kwargs.update(dict(labels=labels, fontsize=fontsize)) #TODO pythonic
        if circles is None: circles = self.arangelon(dlon)
        Basemap.drawmeridians(self, circles, **kwargs)

    def drawgrid(self, lats=None, lons=None, **kwargs):
        """Draw a grid of latitude and longitude lines."""
        self.drawparallels(circles=lats, **kwargs)
        self.drawmeridians(circles=lons, **kwargs)

    def drawdefault(self):
        #TODO accept several kw dicts
        self.drawcoastlines()
        self.fillcontinents(color='coral',lake_color='aqua')
        #self.drawmapboundary(fill_color='aqua')
        self.drawmapboundary()
        self.drawrivers(color='b')

    def draw_bathymetry(self, isobaths):
        """Draw isobaths under water.

        """
        raise NotImplementedError

    def fill_bathymetry(self, cmap):
        """Fill bathymetry under water. (Color plot of sounding depths)

        """
        raise NotImplementedError

    def draw_topography(self, ):
        raise NotImplementedError

    def fill_topography(self, ):
        raise NotImplementedError

    def draw_buoys(self, ):
        """Permanent marker buoys.

        """
        raise NotImplementedError

    def draw_shipping_lanes(self, ):
        """Pretty Self-explanatory.

        """
        raise NotImplementedError

    def draw_track(self, track, **kwargs):
        """Draw a position track on the map (e.g., from Tethys or Daphne).

        """
        raise NotImplementedError
#TODO a track that changes color based on an ancillary variable

    def draw_currents(self, lat, lon, u, v, contourf_magnitude=True, **kwargs):
        """Quiver plot the ocean current.

        """
        if lat.ndim is 2: lat = np.sort(np.unique(lat))
        if lon.ndim is 2: lon = np.sort(np.unique(lon))
        um, vm, xm, ym = self.transform_vector(u, v, lon, lat,
            len(lon), len(lat), returnxy = True, masked = True, order = 1)
        #TODO   pass kwargs to transform_vector?
        #TODO   optional input/output of *m variables above...
        if contourf_magnitude:
            mm = (um**2 + vm**2)**0.5 # 2D current magnitude
            dc = 0.01
#TODO            cmax = max(mm) + dc
            cmax = 0.5 + dc
            mmlevels = np.arange(0, cmax, dc)
            self.contourf(xm, ym, mm, mmlevels, cmap=plt.cm.Blues)
            self.colorbar() #TODO pass kwargs
        q = self.quiver(xm, ym, um, vm, **kwargs)
            #or specify, e.g., width=0.003, scale=400)
        qkey = plt.quiverkey(q, 0.1, -0.2, 1, '1 m/s', labelpos='W')
            #TODO as defaults
        return q, qkey

    def draw_wind(self, wind, altitude, **kwargs):
        """Quiver plot the winds.

        (with or without a mask over the land)
        """
        #TODO redefine args, or unpack wind into u,v,lat,lon
        # transform from spherical to map projection coordinates (rotation
        # and interpolation).
        default = dict(length = 6, linewidth = 0.5,
            barbcolor = 'black', flagcolor = 'cyan', )
        nxv = 25; nyv = 25 #TODO can these be included in defaults?
        default.update(kwargs)
        um, vm, xm, ym = self.transform_vector(u,v,lons,lats,nxv,nyv,returnxy=True)
        self.barbs(xm, ym, um, vm, **default) #TODO what does this return?
        raise NotImplementedError

    def play(self, **kwargs):
        """Generate a movie using inputs with a time dimension.

        """
        raise NotImplementedError

class MontereyBay(Map):
    """A standard map canvas of Monterey Bay.

    """
    def __init__(self, lat_0 = 36.75, lon_0 = -121.0,
        llcrnrlat = 36.5, llcrnrlon = -122.5,
        urcrnrlat = 37, urcrnrlon = -121.75,
        projection = 'tmerc', resolution = 'i', **kwargs):
        """Overloaded method with default args.

        """
        Map.__init__(self, **lib.injectlocals(locals()))
        #TODO some handling if the map is defined in alternate fashion, e.g.:
#m = Basemap(width=920000,height=1100000,
#            resolution='f',projection='tmerc',lon_0=-4.2,lat_0=54.6)

    def drawparallels(self, dlat=.1, **kwargs):
        """Overloaded method with default args.

        """
        Map.drawparallels(self, dlat=dlat, **kwargs)

    def drawmeridians(self, dlon=.1, **kwargs):
        """Overloaded method with default args.

        """
        Map.drawmeridians(self, dlon=dlon, **kwargs)

    def draw_mbari(self, ):
        """Mark location of MBARI on the shore.

        """
        raise NotImplementedError

    def draw_canyon(self, ):
        """Draw boundaries of Monterey Canyon on the map.

        """
        raise NotImplementedError
#TODO other canyons, too?

    def draw_mountains(self, ):
        raise NotImplementedError

    def draw_mars(self,color='red', marker='s', label='MARS', **kwargs):
        """Mark position of MARS cabled observatory.

        """
        plot_kwargs = lib.injectlocals(locals())
        mars_lat = 36 + 42.7481/60 #TODO use proj or write a utility for oa.lib
        mars_lon =  -(122 + 11.2139/60)
        mars_depth = 891 # meters
        x, y = self(mars_lon, mars_lat) # position in projection
        self.plot(x, y, **plot_kwargs)
        #TODO what figure does it go to by default?

<|endoftext|>


<commit_before>import os
import sys
import json
import string
import argparse as ap
from os import path
from glob import glob
from collections import namedtuple

import jedi

global verbose, quiet

def log(msg):
    if verbose:
        sys.stderr.write(msg + '\n')

def error(msg):
    if not quiet:
        sys.stderr.write(msg + '\n')

def main():
    argser = ap.ArgumentParser(description='graph.py is a command that dumps all Python definitions and references found in code rooted at a directory')
    argser.add_argument('dir', help='path to root directory of code')
    argser.add_argument('--pretty', help='pretty print JSON output', action='store_true', default=False)
    argser.add_argument('--verbose', help='verbose', action='store_true', default=False)
    argser.add_argument('--quiet', help='quiet', action='store_true', default=False)

    args = argser.parse_args()
    global verbose, quiet
    verbose, quiet = args.verbose, args.quiet

    if args.dir == '':
        error('target directory must not be empty')
        os.exit(1)

    os.chdir(args.dir)          # set working directory to be source directory

    source_files = glob('**/*.py')
    source_files.extend(glob('*.py'))

    modules_and_files = [(filename_to_module_name(f), f) for f in source_files]
    jedi.api.preload_module([mf[0] for mf in modules_and_files])

    defs = [d for d in get_defs(source_files)]
    refs = [r for r in get_refs(source_files)]

    for module, filename in modules_and_files:
        defs.append(Def(
            Path=module.replace('.', '/'),
            Kind='module',
            Name=string.split(module, '.')[-1],
            File=filename,
            DefStart=0,
            DefEnd=0,
            Exported=True,
            Docstring='',           # TODO: extract module/package-level doc
            Data=None,
        ))

    # De-duplicate definitions (local variables may be defined in more than one
    # place). Could do something smarter here, but for now, just take the first
    # definition that appears.
    unique_defs = []
    unique_def_paths = set([])
    for def_ in defs:
        if not def_.Path in unique_def_paths:
            unique_defs.append(def_)
            unique_def_paths.add(def_.Path)

    json_indent = 2 if args.pretty else None
    print json.dumps({
        'Defs': [d.__dict__ for d in unique_defs],
        'Refs': [r.__dict__ for r in refs],
    }, indent=json_indent)

def get_defs(source_files):
    for source_file in source_files:
        log('getting defs for source file %s' % source_file)
        try:
            source = None
            with open(source_file) as sf:
                source = unicode(sf.read())
            linecoler = LineColToOffConverter(source)

            defs = jedi.api.defined_names(source, path=source_file)
            for def_ in defs:
                for d in get_defs_(def_, source_file, linecoler):
                    yield d
        except Exception as e:
            error('failed to get defs for source file %s: %s' % (source_file, str(e)))

def get_defs_(def_, source_file, linecoler):
    # ignore import definitions because these just redefine things imported from elsewhere
    if def_.type == 'import':
        return

    def__, err = jedi_def_to_def(def_, source_file, linecoler)
    if err is None:
        yield def__
    else:
        error(err)

    if def_.type not in ['function', 'class', 'module']:
        return

    subdefs = def_.defined_names()
    for subdef in subdefs:
        for d in get_defs_(subdef, source_file, linecoler):
            yield d

def jedi_def_to_def(def_, source_file, linecoler):
    full_name, err = full_name_of_def(def_)
    if err is not None:
        return None, err

    start_pos = linecoler.convert(def_.start_pos)
    return Def(
        Path=full_name.replace('.', '/'),
        Kind=def_.type,
        Name=def_.name,
        File=source_file,
        DefStart=start_pos,
        DefEnd=start_pos+len(def_.name),
        Exported=True,          # TODO: not all vars are exported
        Docstring=def_.docstring(),
        Data=None,
    ), None

def get_refs(source_files):
    for source_file in source_files:
        log('getting refs for source file %s' % source_file)
        try:
            parserContext = ParserContext(source_file)
            linecoler = LineColToOffConverter(parserContext.source)
            for name_part, def_ in parserContext.refs():
                try:
                    full_name, err = full_name_of_def(def_, from_ref=True)
                    if err is not None or full_name == '':
                        raise Exception(err)
                    start = linecoler.convert(name_part.start_pos)
                    end = linecoler.convert(name_part.end_pos)
                    yield Ref(
                        DefPath=full_name.replace('.', '/'),
                        DefFile=def_.module_path,
                        Def=False,
                        File=source_file,
                        Start=start,
                        End=end,
                        ToBuiltin=def_.in_builtin_module(),
                    )
                except Exception as e:
                    error('failed to get ref (%s) in source file %s: %s' % (str((name_part, def_)), source_file, str(e)))
        except Exception as e:
            error('failed to get refs for source file %s: %s' % (source_file, str(e)))

def full_name_of_def(def_, from_ref=False):
    # TODO: This function
    # - currently fails for tuple assignments (e.g., 'x, y = 1, 3')
    # - doesn't distinguish between m(module).n(submodule) and m(module).n(contained-variable)

    if def_.in_builtin_module():
        return def_.full_name, None

    full_name = ('%s.%s' % (def_.full_name, def_.name)) if def_.type in set(['statement', 'param']) else def_.full_name

    module_path = def_.module_path
    if from_ref:
        module_path, err = abs_module_path_to_relative_module_path(module_path)
        if err is not None:
            return None, err

    supermodule = supermodule_path(module_path).replace('/', '.')

    # definition definitions' full_name property contains only the promixal module, so we need to add back the parent
    # module components. Luckily, the module_path is relative in this case.
    return path.join(supermodule, full_name), None

def supermodule_path(module_path):
    if path.basename(module_path) == '__init__.py':
        return path.dirname(path.dirname(module_path))
    return path.dirname(module_path)

def abs_module_path_to_relative_module_path(module_path):
    relpath = path.relpath(module_path) # relative from pwd (which is set in main)
    if not relpath.startswith('..'):
        return relpath, None
    components = module_path.split(os.sep)
    pIdx = -1
    for i, cmpt in enumerate(components):
        if cmpt in ['site-packages', 'dist-packages']:
            pIdx = i
            break
    if pIdx != -1:
        return path.join(*components[i+1:]), None

    for i, cmpt in enumerate(components):
        if cmpt.startswith('python'):
            pIdx = i
            break
    if pIdx != -1:
        return path.join(*components[i+1:]), None
    return None, ("could not convert absolute module path %s to relative module path" % module_path)

Def = namedtuple('Def', ['Path', 'Kind', 'Name', 'File', 'DefStart', 'DefEnd', 'Exported', 'Docstring', 'Data'])
Ref = namedtuple('Ref', ['DefPath', 'DefFile', 'Def', 'File', 'Start', 'End', "ToBuiltin"])

class ParserContext(object):
    def __init__(self, source_file):
        self.source_file = source_file
        with open(source_file) as sf:
            self.source = unicode(sf.read())
            self.parser = jedi.parser.Parser(self.source, source_file)

    def refs(self):
        for r in self.scope_refs(self.parser.module):
            yield r

    def scope_refs(self, scope):
        for import_ in scope.imports:
            for r in self.import_refs(import_):
                yield r
        for stmt in scope.statements:
            for r in self.stmt_refs(stmt):
                yield r
        for subscope in scope.subscopes:
            for r in self.scope_refs(subscope):
                yield r

    def import_refs(self, import_):
        for name in import_.get_all_import_names():
            for name_part in name.names:
                defs = jedi.api.Script(
                    path=self.source_file,
                    line=name_part.start_pos[0],
                    column=name_part.start_pos[1],
                ).goto_assignments()
                for def_ in defs:
                    yield (name_part, def_)

    def stmt_refs(self, stmt):
        if isinstance(stmt, jedi.parser.representation.KeywordStatement):
            return
        if isinstance(stmt, jedi.parser.representation.Flow):
            return

        for token in stmt._token_list:
            if not isinstance(token, jedi.parser.representation.Name):
                continue
            for name_part in token.names:
                # Note: we call goto_definitions instead of goto_assignments,
                # because otherwise the reference will not follow imports (and
                # also generates bogus local definitions whose paths conflict
                # with those of actual definitions). This uses a modified
                # goto_definitions (resolve_variables_to_types option) that
                # *DOES NOT* follow assignment statements to resolve variables
                # to types (because that's not what we want).
                defs = jedi.api.Script(
                    path=self.source_file,
                    line=name_part.start_pos[0],
                    column=name_part.start_pos[1],
                    resolve_variables_to_types=False,
                ).goto_definitions()

                # Note(beyang): For now, only yield the first definition.
                # Otherwise, multiple references to multiple definitions will
                # yield dup references. In the future, might want to do
                # something smarter here.
                i = 0
                for def_ in defs:
                    if i > 0: break
                    yield (name_part, def_)
                    i += 1

def resolve_import_paths(scopes):
    for s in scopes.copy():
        if isinstance(s, jedi.evaluate.imports.ImportWrapper):
            scopes.remove(s)
            scopes.update(resolve_import_paths(set(s.follow())))
    return scopes

def filename_to_module_name(filename):
    if path.basename(filename) == '__init__.py':
        return path.dirname(filename).replace('/', '.')
    return path.splitext(filename)[0].replace('/', '.')

class LineColToOffConverter(object):
    def __init__(self, source):
        source_lines = source.split('\n')
        cumulative_off = [0]
        for line in source_lines:
            cumulative_off.append(cumulative_off[-1] + len(line) + 1)
        self._cumulative_off = cumulative_off

    # Converts from (line, col) position to byte offset. line is 1-indexed, col is 0-indexed
    def convert(self, linecol):
        line, col = linecol[0] - 1, linecol[1]         # convert line to 0-indexed
        if line >= len(self._cumulative_off):
            return None, 'requested line out of bounds %d > %d' % (line+1, len(self._cumulative_off)-1)
        return self._cumulative_off[line] + col

if __name__ == '__main__':
    main()
<commit_msg>parse return statements<commit_after>import os
import sys
import json
import string
import argparse as ap
from os import path
from glob import glob
from collections import namedtuple

import jedi

global verbose, quiet

def log(msg):
    if verbose:
        sys.stderr.write(msg + '\n')

def error(msg):
    if not quiet:
        sys.stderr.write(msg + '\n')

def main():
    argser = ap.ArgumentParser(description='graph.py is a command that dumps all Python definitions and references found in code rooted at a directory')
    argser.add_argument('dir', help='path to root directory of code')
    argser.add_argument('--pretty', help='pretty print JSON output', action='store_true', default=False)
    argser.add_argument('--verbose', help='verbose', action='store_true', default=False)
    argser.add_argument('--quiet', help='quiet', action='store_true', default=False)

    args = argser.parse_args()
    global verbose, quiet
    verbose, quiet = args.verbose, args.quiet

    if args.dir == '':
        error('target directory must not be empty')
        os.exit(1)

    os.chdir(args.dir)          # set working directory to be source directory

    source_files = glob('**/*.py')
    source_files.extend(glob('*.py'))

    modules_and_files = [(filename_to_module_name(f), f) for f in source_files]
    jedi.api.preload_module([mf[0] for mf in modules_and_files])

    defs = [d for d in get_defs(source_files)]
    refs = [r for r in get_refs(source_files)]

    for module, filename in modules_and_files:
        defs.append(Def(
            Path=module.replace('.', '/'),
            Kind='module',
            Name=string.split(module, '.')[-1],
            File=filename,
            DefStart=0,
            DefEnd=0,
            Exported=True,
            Docstring='',           # TODO: extract module/package-level doc
            Data=None,
        ))

    # De-duplicate definitions (local variables may be defined in more than one
    # place). Could do something smarter here, but for now, just take the first
    # definition that appears.
    unique_defs = []
    unique_def_paths = set([])
    for def_ in defs:
        if not def_.Path in unique_def_paths:
            unique_defs.append(def_)
            unique_def_paths.add(def_.Path)

    json_indent = 2 if args.pretty else None
    print json.dumps({
        'Defs': [d.__dict__ for d in unique_defs],
        'Refs': [r.__dict__ for r in refs],
    }, indent=json_indent)

def get_defs(source_files):
    for source_file in source_files:
        log('getting defs for source file %s' % source_file)
        try:
            source = None
            with open(source_file) as sf:
                source = unicode(sf.read())
            linecoler = LineColToOffConverter(source)

            defs = jedi.api.defined_names(source, path=source_file)
            for def_ in defs:
                for d in get_defs_(def_, source_file, linecoler):
                    yield d
        except Exception as e:
            error('failed to get defs for source file %s: %s' % (source_file, str(e)))

def get_defs_(def_, source_file, linecoler):
    # ignore import definitions because these just redefine things imported from elsewhere
    if def_.type == 'import':
        return

    def__, err = jedi_def_to_def(def_, source_file, linecoler)
    if err is None:
        yield def__
    else:
        error(err)

    if def_.type not in ['function', 'class', 'module']:
        return

    subdefs = def_.defined_names()
    for subdef in subdefs:
        for d in get_defs_(subdef, source_file, linecoler):
            yield d

def jedi_def_to_def(def_, source_file, linecoler):
    full_name, err = full_name_of_def(def_)
    if err is not None:
        return None, err

    start_pos = linecoler.convert(def_.start_pos)
    return Def(
        Path=full_name.replace('.', '/'),
        Kind=def_.type,
        Name=def_.name,
        File=source_file,
        DefStart=start_pos,
        DefEnd=start_pos+len(def_.name),
        Exported=True,          # TODO: not all vars are exported
        Docstring=def_.docstring(),
        Data=None,
    ), None

def get_refs(source_files):
    for source_file in source_files:
        log('getting refs for source file %s' % source_file)
        try:
            parserContext = ParserContext(source_file)
            linecoler = LineColToOffConverter(parserContext.source)
            for name_part, def_ in parserContext.refs():
                try:
                    full_name, err = full_name_of_def(def_, from_ref=True)
                    if err is not None or full_name == '':
                        raise Exception(err)
                    start = linecoler.convert(name_part.start_pos)
                    end = linecoler.convert(name_part.end_pos)
                    yield Ref(
                        DefPath=full_name.replace('.', '/'),
                        DefFile=def_.module_path,
                        Def=False,
                        File=source_file,
                        Start=start,
                        End=end,
                        ToBuiltin=def_.in_builtin_module(),
                    )
                except Exception as e:
                    error('failed to get ref (%s) in source file %s: %s' % (str((name_part, def_)), source_file, str(e)))
        except Exception as e:
            error('failed to get refs for source file %s: %s' % (source_file, str(e)))

def full_name_of_def(def_, from_ref=False):
    # TODO: This function
    # - currently fails for tuple assignments (e.g., 'x, y = 1, 3')
    # - doesn't distinguish between m(module).n(submodule) and m(module).n(contained-variable)

    if def_.in_builtin_module():
        return def_.full_name, None

    full_name = ('%s.%s' % (def_.full_name, def_.name)) if def_.type in set(['statement', 'param']) else def_.full_name

    module_path = def_.module_path
    if from_ref:
        module_path, err = abs_module_path_to_relative_module_path(module_path)
        if err is not None:
            return None, err

    supermodule = supermodule_path(module_path).replace('/', '.')

    # definition definitions' full_name property contains only the promixal module, so we need to add back the parent
    # module components. Luckily, the module_path is relative in this case.
    return path.join(supermodule, full_name), None

def supermodule_path(module_path):
    if path.basename(module_path) == '__init__.py':
        return path.dirname(path.dirname(module_path))
    return path.dirname(module_path)

def abs_module_path_to_relative_module_path(module_path):
    relpath = path.relpath(module_path) # relative from pwd (which is set in main)
    if not relpath.startswith('..'):
        return relpath, None
    components = module_path.split(os.sep)
    pIdx = -1
    for i, cmpt in enumerate(components):
        if cmpt in ['site-packages', 'dist-packages']:
            pIdx = i
            break
    if pIdx != -1:
        return path.join(*components[i+1:]), None

    for i, cmpt in enumerate(components):
        if cmpt.startswith('python'):
            pIdx = i
            break
    if pIdx != -1:
        return path.join(*components[i+1:]), None
    return None, ("could not convert absolute module path %s to relative module path" % module_path)

Def = namedtuple('Def', ['Path', 'Kind', 'Name', 'File', 'DefStart', 'DefEnd', 'Exported', 'Docstring', 'Data'])
Ref = namedtuple('Ref', ['DefPath', 'DefFile', 'Def', 'File', 'Start', 'End', "ToBuiltin"])

class ParserContext(object):
    def __init__(self, source_file):
        self.source_file = source_file
        with open(source_file) as sf:
            self.source = unicode(sf.read())
            self.parser = jedi.parser.Parser(self.source, source_file)

    def refs(self):
        for r in self.scope_refs(self.parser.module):
            yield r

    def scope_refs(self, scope):
        for import_ in scope.imports:
            for r in self.import_refs(import_):
                yield r
        for stmt in scope.statements:
            for r in self.stmt_refs(stmt):
                yield r
        for ret in scope.returns:
            for r in self.stmt_refs(ret):
                yield r
        for subscope in scope.subscopes:
            for r in self.scope_refs(subscope):
                yield r

    def import_refs(self, import_):
        for name in import_.get_all_import_names():
            for name_part in name.names:
                defs = jedi.api.Script(
                    path=self.source_file,
                    line=name_part.start_pos[0],
                    column=name_part.start_pos[1],
                ).goto_assignments()
                for def_ in defs:
                    yield (name_part, def_)

    def stmt_refs(self, stmt):
        if isinstance(stmt, jedi.parser.representation.KeywordStatement):
            return
        if isinstance(stmt, jedi.parser.representation.Flow):
            return

        for token in stmt._token_list:
            if not isinstance(token, jedi.parser.representation.Name):
                continue
            for name_part in token.names:
                # Note: we call goto_definitions instead of goto_assignments,
                # because otherwise the reference will not follow imports (and
                # also generates bogus local definitions whose paths conflict
                # with those of actual definitions). This uses a modified
                # goto_definitions (resolve_variables_to_types option) that
                # *DOES NOT* follow assignment statements to resolve variables
                # to types (because that's not what we want).
                defs = jedi.api.Script(
                    path=self.source_file,
                    line=name_part.start_pos[0],
                    column=name_part.start_pos[1],
                    resolve_variables_to_types=False,
                ).goto_definitions()

                # Note(beyang): For now, only yield the first definition.
                # Otherwise, multiple references to multiple definitions will
                # yield dup references. In the future, might want to do
                # something smarter here.
                i = 0
                for def_ in defs:
                    if i > 0: break
                    yield (name_part, def_)
                    i += 1

def resolve_import_paths(scopes):
    for s in scopes.copy():
        if isinstance(s, jedi.evaluate.imports.ImportWrapper):
            scopes.remove(s)
            scopes.update(resolve_import_paths(set(s.follow())))
    return scopes

def filename_to_module_name(filename):
    if path.basename(filename) == '__init__.py':
        return path.dirname(filename).replace('/', '.')
    return path.splitext(filename)[0].replace('/', '.')

class LineColToOffConverter(object):
    def __init__(self, source):
        source_lines = source.split('\n')
        cumulative_off = [0]
        for line in source_lines:
            cumulative_off.append(cumulative_off[-1] + len(line) + 1)
        self._cumulative_off = cumulative_off

    # Converts from (line, col) position to byte offset. line is 1-indexed, col is 0-indexed
    def convert(self, linecol):
        line, col = linecol[0] - 1, linecol[1]         # convert line to 0-indexed
        if line >= len(self._cumulative_off):
            return None, 'requested line out of bounds %d > %d' % (line+1, len(self._cumulative_off)-1)
        return self._cumulative_off[line] + col

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


<commit_before>"""
Generate Netkit configuration files for a network
"""
from mako.lookup import TemplateLookup

# TODO: merge these imports but make consistent across compilers
from pkg_resources import resource_filename
import pkg_resources

import os

#import network as network

import logging
LOG = logging.getLogger("ANK")

import shutil
import glob
import itertools

import AutoNetkit as ank
from AutoNetkit import config

import pprint
pp = pprint.PrettyPrinter(indent=4)
import tarfile
import time

# Check can write to template cache directory
#TODO: make function to provide cache directory
#TODO: move this into config
template_cache_dir = config.template_cache_dir

template_dir =  resource_filename("AutoNetkit","lib/templates")
lookup = TemplateLookup(directories=[ template_dir ],
        module_directory= template_cache_dir,
        #cache_type='memory',
        #cache_enabled=True,
        )

def lab_dir():
    """Lab directory for junos configs"""
    return config.junos_dir

def router_conf_dir():
    """Directory for individual Junos router configs"""
    return os.path.join(lab_dir(), "configset")

def router_conf_file(network, node):
    """Returns filename for config file for router"""
    return "%s.conf" % ank.rtr_folder_name(network, node)

def router_conf_path(network, node):
    """ Returns full path to router config file"""
    r_file = router_conf_file(network, node)
    return os.path.join(router_conf_dir(), r_file)

def int_id_em(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg em1"""
# Junosphere uses em0 for external link
    numeric_id += 1
    return 'em%s' % numeric_id

def int_id_ge(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg ge-0/0/1"""
# Junosphere uses ge/0/0/0 for external link
    numeric_id += 1
    return 'ge-0/0/%s' % numeric_id

def logical_int_id_ge(numeric_id):
    """ For routing protocols, refer to logical int id:
    ge-0/0/1 becomes ge-0/0/1.0"""
    return int_id_ge(numeric_id) + ".0"

class JunosCompiler:
    """Compiler main"""

    def __init__(self, network, services, igp):
        self.network = network
        self.services = services
        self.igp = igp

    def initialise(self):
        """Creates lab folder structure"""
        if not os.path.isdir(lab_dir()):
            os.mkdir(lab_dir())
        else:
            for item in glob.iglob(os.path.join(lab_dir(), "*")):
                if os.path.isdir(item):
                    shutil.rmtree(item)
                else:
                    os.unlink(item)

        # Directory to put config files into
        if not os.path.isdir(router_conf_dir()):
            os.mkdir(router_conf_dir())
        return

    def configure_junosphere(self):
        """Configure Junosphere topology structure"""
        vmm_template = lookup.get_template("junos/topology_vmm.mako")
        topology_data = {}
        # Generator for private0, private1, etc
        #TODO: replace count with xrange up to limit of interfaces, eg 64 in Junosphere
        bridge_id_generator = ('private%s'%i for i in itertools.count(0))
        collision_to_bridge_mapping = {}

        for node in self.network.graph:
            hostname = ank.fqdn(self.network, node)
            topology_data[hostname] = {
                    'image': 'VJX1000_LATEST',
                    'config': router_conf_file(self.network, node),
                    'interfaces': [],
                    }
            for src, dst, data in self.network.graph.edges(node, data=True):
                print "data", data
                subnet = data['sn']
                description = 'Interface %s -> %s' % (
                        ank.fqdn(self.network, src),
                        ank.fqdn(self.network, dst))
# Bridge information for topology config
                if subnet in collision_to_bridge_mapping:
# Use bridge allocated for this subnet
                    bridge_id = collision_to_bridge_mapping[subnet]
                else:
# Allocate a bridge for this subnet
                    bridge_id = bridge_id_generator.next()
                    collision_to_bridge_mapping[subnet] = bridge_id

                topology_data[hostname]['interfaces'].append({
                    'description': description,
                    'id': int_id_em(data['id']),
                    'id_ge':  int_id_ge(data['id']),
                    'bridge_id': bridge_id,
                    })

        if len(collision_to_bridge_mapping) > 64:
            LOG.warn("AutoNetkit does not currently support more"
                    " than 123 network links for Junosphere")

        vmm_file = os.path.join(lab_dir(), "topology.vmm")
        with open( vmm_file, 'w') as f_vmm:
            f_vmm.write( vmm_template.render(
                topology_data = topology_data,
                ))

    def configure_interfaces(self, node):
        """Interface configuration"""
        lo_ip = self.network.lo_ip(node)
        interfaces = []

        interfaces.append({
            'id':          'lo0',
            'ip':           str(lo_ip.ip),
            'netmask':      str(lo_ip.netmask),
            'prefixlen':    str(lo_ip.prefixlen),
            'net_ent_title': ank.ip_to_net_ent_title(lo_ip),
            'description': 'Loopback',
        })

        # Add em0.0 for Qemu
#TODO: make this enabled with switch for QEMU
        """
        interfaces.append({
            'id':          'em0',# modified from "em0.0"
            'ip':           str(self.network[node].get('tap_ip')),
            'netmask':      str(tap_subnet.netmask),
            'prefixlen':    32,# modified into 32
            'description': 'Admin for Qemu',
        })
        """

        for src, dst, data in self.network.graph.edges(node, data=True):
            subnet = data['sn']
            int_id = int_id_ge(data['id'])
            description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))

# Interface information for router config
            interfaces.append({
                'id':          int_id,
                'ip':           str(data['ip']),
                'prefixlen':    str(subnet.prefixlen),
                'broadcast':    str(subnet.broadcast),
                'description':  description,
            })

        return interfaces

    def configure_igp(self, node, igp_graph):
        """igp configuration"""
        default_weight = 1
        igp_interfaces = []
        if igp_graph.degree(node) > 0:
            # Only start IGP process if IGP links
            igp_interfaces.append({ 'id': 'lo0', 'passive': True})
            for src, dst, data in igp_graph.edges(node, data=True):
                int_id = logical_int_id_ge(data['id'])
                description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))
                igp_interfaces.append({
                    'id':       int_id,
                    'weight':   data.get('weight', default_weight),
                    'description': description,
                    })
        return igp_interfaces

    def configure_bgp(self, node, physical_graph, ibgp_graph, ebgp_graph):
        """ BGP configuration"""
        if len(ebgp_graph.edges()) == 0:
# Don't configure iBGP or eBGP if no eBGP edges
            LOG.debug("Skipping BGP configuration for %s as no eBGP edges" % node)
            return

#TODO: put comments in for junos bgp peerings

        bgp_groups = {}
        if node in ibgp_graph:
            internal_peers = []
            for peer in ibgp_graph.neighbors(node):
                internal_peers.append({'id': self.network.lo_ip(peer).ip})
            bgp_groups['internal_peers'] = {
                    'type': 'internal',
                    'neighbors': internal_peers
                    }

        ibgp_neighbor_list = []
        ibgp_rr_client_list = []
        if node in ibgp_graph:
            for src, neigh, data in ibgp_graph.edges(node, data=True):
                description = data.get("rr_dir") + " to " + ank.fqdn(self.network, neigh)
                if data.get('rr_dir') == 'down':
                    ibgp_rr_client_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })
                elif (data.get('rr_dir') in set(['up', 'over', 'peer'])
                        or data.get('rr_dir') is None):
                    ibgp_neighbor_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })

        bgp_groups['internal_peers'] = {
            'type': 'internal',
            'neighbors': ibgp_neighbor_list
            }
        if len(ibgp_rr_client_list):
            bgp_groups['internal_rr'] = {
                    'type': 'internal',
                    'neighbors': ibgp_rr_client_list,
                    'cluster': self.network.lo_ip(node).ip,
                    }

        if node in ebgp_graph:
            external_peers = []
            for peer in ebgp_graph.neighbors(node):
                peer_ip = physical_graph[peer][node]['ip']
                external_peers.append({
                    'id': peer_ip,
                    'peer_as': self.network.asn(peer)})
            bgp_groups['external_peers'] = {
                    'type': 'external',
                    'neighbors': external_peers}

        return bgp_groups

    def configure_junos(self):
        """ Configures Junos"""
        LOG.info("Configuring Junos")
        junos_template = lookup.get_template("junos/junos.mako")
        ank_version = pkg_resources.get_distribution("AutoNetkit").version

        physical_graph = self.network.graph
        igp_graph = ank.igp_graph(self.network)
        ibgp_graph = ank.get_ibgp_graph(self.network)
        ebgp_graph = ank.get_ebgp_graph(self.network)

        #TODO: correct this router type selector
        for node in self.network.q(platform="NETKIT"):
            asn = self.network.asn(node)
            network_list = []
            lo_ip = self.network.lo_ip(node)

            interfaces = self.configure_interfaces(node)
            igp_interfaces = self.configure_igp(node, igp_graph)
            bgp_groups = self.configure_bgp(node, physical_graph, ibgp_graph, ebgp_graph)

            # advertise AS subnet
            adv_subnet = self.network.ip_as_allocs[asn]
            if not adv_subnet in network_list:
                network_list.append(adv_subnet)

            juniper_filename = router_conf_path(self.network, node)
            with open( juniper_filename, 'w') as f_jun:
                f_jun.write( junos_template.render(
                    hostname = ank.fqdn(self.network, node),
                    username = 'autonetkit',
                    interfaces=interfaces,
                    igp_interfaces=igp_interfaces,
                    igp_protocol = self.igp,
                    asn = asn,
                    lo_ip=lo_ip,
                    router_id = lo_ip.ip,
                    network_list = network_list,
                    bgp_groups = bgp_groups,
                    ank_version = ank_version,
                    ))

    def configure(self):
        self.configure_junosphere()
        self.configure_junos()
# create .tgz
        tar_filename = "junos_%s.tar.gz" % time.strftime("%Y%m%d_%H%M",
                time.localtime())
        tar = tarfile.open(os.path.join(config.ank_main_dir,
            tar_filename), "w:gz")
# arcname to flatten file structure
        tar.add(lab_dir(), arcname="")
        self.network.compiled_labs['junos'] = tar_filename
        tar.close()
<commit_msg>remove debug output<commit_after>"""
Generate Netkit configuration files for a network
"""
from mako.lookup import TemplateLookup

# TODO: merge these imports but make consistent across compilers
from pkg_resources import resource_filename
import pkg_resources

import os

#import network as network

import logging
LOG = logging.getLogger("ANK")

import shutil
import glob
import itertools

import AutoNetkit as ank
from AutoNetkit import config

import pprint
pp = pprint.PrettyPrinter(indent=4)
import tarfile
import time

# Check can write to template cache directory
#TODO: make function to provide cache directory
#TODO: move this into config
template_cache_dir = config.template_cache_dir

template_dir =  resource_filename("AutoNetkit","lib/templates")
lookup = TemplateLookup(directories=[ template_dir ],
        module_directory= template_cache_dir,
        #cache_type='memory',
        #cache_enabled=True,
        )

def lab_dir():
    """Lab directory for junos configs"""
    return config.junos_dir

def router_conf_dir():
    """Directory for individual Junos router configs"""
    return os.path.join(lab_dir(), "configset")

def router_conf_file(network, node):
    """Returns filename for config file for router"""
    return "%s.conf" % ank.rtr_folder_name(network, node)

def router_conf_path(network, node):
    """ Returns full path to router config file"""
    r_file = router_conf_file(network, node)
    return os.path.join(router_conf_dir(), r_file)

def int_id_em(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg em1"""
# Junosphere uses em0 for external link
    numeric_id += 1
    return 'em%s' % numeric_id

def int_id_ge(numeric_id):
    """Returns Junos format interface ID for an AutoNetkit interface ID
    eg ge-0/0/1"""
# Junosphere uses ge/0/0/0 for external link
    numeric_id += 1
    return 'ge-0/0/%s' % numeric_id

def logical_int_id_ge(numeric_id):
    """ For routing protocols, refer to logical int id:
    ge-0/0/1 becomes ge-0/0/1.0"""
    return int_id_ge(numeric_id) + ".0"

class JunosCompiler:
    """Compiler main"""

    def __init__(self, network, services, igp):
        self.network = network
        self.services = services
        self.igp = igp

    def initialise(self):
        """Creates lab folder structure"""
        if not os.path.isdir(lab_dir()):
            os.mkdir(lab_dir())
        else:
            for item in glob.iglob(os.path.join(lab_dir(), "*")):
                if os.path.isdir(item):
                    shutil.rmtree(item)
                else:
                    os.unlink(item)

        # Directory to put config files into
        if not os.path.isdir(router_conf_dir()):
            os.mkdir(router_conf_dir())
        return

    def configure_junosphere(self):
        """Configure Junosphere topology structure"""
        vmm_template = lookup.get_template("junos/topology_vmm.mako")
        topology_data = {}
        # Generator for private0, private1, etc
        #TODO: replace count with xrange up to limit of interfaces, eg 64 in Junosphere
        bridge_id_generator = ('private%s'%i for i in itertools.count(0))
        collision_to_bridge_mapping = {}

        for node in self.network.graph:
            hostname = ank.fqdn(self.network, node)
            topology_data[hostname] = {
                    'image': 'VJX1000_LATEST',
                    'config': router_conf_file(self.network, node),
                    'interfaces': [],
                    }
            for src, dst, data in self.network.graph.edges(node, data=True):
                subnet = data['sn']
                description = 'Interface %s -> %s' % (
                        ank.fqdn(self.network, src),
                        ank.fqdn(self.network, dst))
# Bridge information for topology config
                if subnet in collision_to_bridge_mapping:
# Use bridge allocated for this subnet
                    bridge_id = collision_to_bridge_mapping[subnet]
                else:
# Allocate a bridge for this subnet
                    bridge_id = bridge_id_generator.next()
                    collision_to_bridge_mapping[subnet] = bridge_id

                topology_data[hostname]['interfaces'].append({
                    'description': description,
                    'id': int_id_em(data['id']),
                    'id_ge':  int_id_ge(data['id']),
                    'bridge_id': bridge_id,
                    })

        if len(collision_to_bridge_mapping) > 64:
            LOG.warn("AutoNetkit does not currently support more"
                    " than 123 network links for Junosphere")

        vmm_file = os.path.join(lab_dir(), "topology.vmm")
        with open( vmm_file, 'w') as f_vmm:
            f_vmm.write( vmm_template.render(
                topology_data = topology_data,
                ))

    def configure_interfaces(self, node):
        """Interface configuration"""
        lo_ip = self.network.lo_ip(node)
        interfaces = []

        interfaces.append({
            'id':          'lo0',
            'ip':           str(lo_ip.ip),
            'netmask':      str(lo_ip.netmask),
            'prefixlen':    str(lo_ip.prefixlen),
            'net_ent_title': ank.ip_to_net_ent_title(lo_ip),
            'description': 'Loopback',
        })

        # Add em0.0 for Qemu
#TODO: make this enabled with switch for QEMU
        """
        interfaces.append({
            'id':          'em0',# modified from "em0.0"
            'ip':           str(self.network[node].get('tap_ip')),
            'netmask':      str(tap_subnet.netmask),
            'prefixlen':    32,# modified into 32
            'description': 'Admin for Qemu',
        })
        """

        for src, dst, data in self.network.graph.edges(node, data=True):
            subnet = data['sn']
            int_id = int_id_ge(data['id'])
            description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))

# Interface information for router config
            interfaces.append({
                'id':          int_id,
                'ip':           str(data['ip']),
                'prefixlen':    str(subnet.prefixlen),
                'broadcast':    str(subnet.broadcast),
                'description':  description,
            })

        return interfaces

    def configure_igp(self, node, igp_graph):
        """igp configuration"""
        default_weight = 1
        igp_interfaces = []
        if igp_graph.degree(node) > 0:
            # Only start IGP process if IGP links
            igp_interfaces.append({ 'id': 'lo0', 'passive': True})
            for src, dst, data in igp_graph.edges(node, data=True):
                int_id = logical_int_id_ge(data['id'])
                description = 'Interface %s -> %s' % (
                    ank.fqdn(self.network, src),
                    ank.fqdn(self.network, dst))
                igp_interfaces.append({
                    'id':       int_id,
                    'weight':   data.get('weight', default_weight),
                    'description': description,
                    })
        return igp_interfaces

    def configure_bgp(self, node, physical_graph, ibgp_graph, ebgp_graph):
        """ BGP configuration"""
        if len(ebgp_graph.edges()) == 0:
# Don't configure iBGP or eBGP if no eBGP edges
            LOG.debug("Skipping BGP configuration for %s as no eBGP edges" % node)
            return

#TODO: put comments in for junos bgp peerings

        bgp_groups = {}
        if node in ibgp_graph:
            internal_peers = []
            for peer in ibgp_graph.neighbors(node):
                internal_peers.append({'id': self.network.lo_ip(peer).ip})
            bgp_groups['internal_peers'] = {
                    'type': 'internal',
                    'neighbors': internal_peers
                    }

        ibgp_neighbor_list = []
        ibgp_rr_client_list = []
        if node in ibgp_graph:
            for src, neigh, data in ibgp_graph.edges(node, data=True):
                description = data.get("rr_dir") + " to " + ank.fqdn(self.network, neigh)
                if data.get('rr_dir') == 'down':
                    ibgp_rr_client_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })
                elif (data.get('rr_dir') in set(['up', 'over', 'peer'])
                        or data.get('rr_dir') is None):
                    ibgp_neighbor_list.append(
                            {
                                'id':  self.network.lo_ip(neigh).ip,
                                'description':      description,
                                })

        bgp_groups['internal_peers'] = {
            'type': 'internal',
            'neighbors': ibgp_neighbor_list
            }
        if len(ibgp_rr_client_list):
            bgp_groups['internal_rr'] = {
                    'type': 'internal',
                    'neighbors': ibgp_rr_client_list,
                    'cluster': self.network.lo_ip(node).ip,
                    }

        if node in ebgp_graph:
            external_peers = []
            for peer in ebgp_graph.neighbors(node):
                peer_ip = physical_graph[peer][node]['ip']
                external_peers.append({
                    'id': peer_ip,
                    'peer_as': self.network.asn(peer)})
            bgp_groups['external_peers'] = {
                    'type': 'external',
                    'neighbors': external_peers}

        return bgp_groups

    def configure_junos(self):
        """ Configures Junos"""
        LOG.info("Configuring Junos")
        junos_template = lookup.get_template("junos/junos.mako")
        ank_version = pkg_resources.get_distribution("AutoNetkit").version

        physical_graph = self.network.graph
        igp_graph = ank.igp_graph(self.network)
        ibgp_graph = ank.get_ibgp_graph(self.network)
        ebgp_graph = ank.get_ebgp_graph(self.network)

        #TODO: correct this router type selector
        for node in self.network.q(platform="NETKIT"):
            asn = self.network.asn(node)
            network_list = []
            lo_ip = self.network.lo_ip(node)

            interfaces = self.configure_interfaces(node)
            igp_interfaces = self.configure_igp(node, igp_graph)
            bgp_groups = self.configure_bgp(node, physical_graph, ibgp_graph, ebgp_graph)

            # advertise AS subnet
            adv_subnet = self.network.ip_as_allocs[asn]
            if not adv_subnet in network_list:
                network_list.append(adv_subnet)

            juniper_filename = router_conf_path(self.network, node)
            with open( juniper_filename, 'w') as f_jun:
                f_jun.write( junos_template.render(
                    hostname = ank.fqdn(self.network, node),
                    username = 'autonetkit',
                    interfaces=interfaces,
                    igp_interfaces=igp_interfaces,
                    igp_protocol = self.igp,
                    asn = asn,
                    lo_ip=lo_ip,
                    router_id = lo_ip.ip,
                    network_list = network_list,
                    bgp_groups = bgp_groups,
                    ank_version = ank_version,
                    ))

    def configure(self):
        self.configure_junosphere()
        self.configure_junos()
# create .tgz
        tar_filename = "junos_%s.tar.gz" % time.strftime("%Y%m%d_%H%M",
                time.localtime())
        tar = tarfile.open(os.path.join(config.ank_main_dir,
            tar_filename), "w:gz")
# arcname to flatten file structure
        tar.add(lab_dir(), arcname="")
        self.network.compiled_labs['junos'] = tar_filename
        tar.close()
<|endoftext|>


<commit_before>from collections import namedtuple

from . coretypes import MTProtoKeyNotReadyError
from . crypt import SHA1, RSA2048Key, AES256Key, ige_encrypt
from . util import substr

def _aes_encrypt(data, aes_key, aes_iv):
    return ige_encrypt(data, aes_key, aes_iv)

def aes_encrypt(data, auth_key, msg_key):
    aes_key, aes_iv = auth_key.prepare_aes(msg_key)
    return _aes_encrypt(data, aes_key, aes_iv)

def aes_encrypt_local(data, auth_key, msg_key):
    aes_key, aes_iv = auth_key.prepare_aes(msg_key, False)
    return _aes_encrypt(data, aes_key, aes_iv)

class MTProtoAuthKey:

    def __init__(self):
        self._dc = None
        self._key_id = None
        self._key = None

    def set_key(self, key):
        self._key = RSA2048Key(key)
        self._key_id = SHA1(self._key)[-8:]

    def created(self):
        return self._key_id is not None

    @property
    def dc(self):
        return self._dc

    @dc.setter
    def dc(self, dc):
        self._dc = dc

    @property
    def key_id(self):
        if not self.created():
            raise MTProtoKeyNotReadyError("key_id")

        return self._key_id

    def prepare_aes(self, msg_key, send=True):
        """
        Defining AES Key and Initialization Vector

        The 2048-bit authorization key (auth_key) and the 128-bit message key (msg_key) are used to compute a
        56-bit AES key (aes_key) and a 256-bit initialization vector (aes_iv) which are subsequently used to
        encrypt the part of the message to be encrypted (i. e. everything with the exception of the external
            header which is added later) with AES-256 in infinite garble extension (IGE) mode.

        The algorithm for computing aes_key and aes_iv from auth_key and msg_key is as follows:

            sha1_a = SHA1 (msg_key + substr (auth_key, x, 32));
            sha1_b = SHA1 (substr (auth_key, 32+x, 16) + msg_key + substr (auth_key, 48+x, 16));
            sha1_с = SHA1 (substr (auth_key, 64+x, 32) + msg_key);
            sha1_d = SHA1 (msg_key + substr (auth_key, 96+x, 32));
            aes_key = substr (sha1_a, 0, 8) + substr (sha1_b, 8, 12) + substr (sha1_c, 4, 12);
            aes_iv = substr (sha1_a, 8, 12) + substr (sha1_b, 0, 8) + substr (sha1_c, 16, 4) + substr (sha1_d, 0, 8);

        where x = 0 for messages from client to server and x = 8 for those from server to client.

        The lower-order 1024 bits of auth_key are not involved in the computation. They may (together with the
        remaining bits or separately) be used on the client device to encrypt the local copy of the data
        received from the server. The 512 lower-order bits of auth_key are not stored on the server;
        therefore, if the client device uses them to encrypt local data and the user loses the key or the
        password, data decryption of local data is impossible (even if data from the server could be obtained).

        When AES is used to encrypt a block of data of a length not divisible by 16 bytes, the data is padded
        with random bytes to the smallest length divisible by 16 bytes immediately prior to being encrypted.
        """

        auth_key = self._key

        x = 8 if send else 0

        sha1_a = SHA1(msg_key + substr(auth_key, x, 32))
        sha1_b = SHA1(substr(auth_key, 32+x, 16) + msg_key + substr(auth_key, 48+x, 16))
        sha1_с = SHA1(substr(auth_key, 64+x, 32) + msg_key)
        sha1_d = SHA1(msg_key + substr(auth_key, 96+x, 32))
        aes_key = substr(sha1_a, 0, 8) + substr(sha1_b, 8, 12) + substr(sha1_с, 4, 12)
        aes_iv = substr(sha1_a, 8, 12) + substr(sha1_b, 0, 8) + substr(sha1_с, 16, 4) + substr(sha1_d, 0, 8)

        return AES256Key(aes_key), AES256Key(aes_iv)
<commit_msg>Now recieving very long response message. Possibly valid, need to verify<commit_after>from collections import namedtuple

from . coretypes import MTProtoKeyNotReadyError
from . crypt import SHA1, RSA2048Key, AES256Key, ige_encrypt
from . util import substr

def _aes_encrypt(data, aes_key, aes_iv):
    return ige_encrypt(data, aes_key, aes_iv)

def aes_encrypt(data, auth_key, msg_key):
    aes_key, aes_iv = auth_key.prepare_aes(msg_key)
    return _aes_encrypt(data, aes_key, aes_iv)

def aes_encrypt_local(data, auth_key, msg_key):
    aes_key, aes_iv = auth_key.prepare_aes(msg_key, False)
    return _aes_encrypt(data, aes_key, aes_iv)

class MTProtoAuthKey:

    def __init__(self):
        self._dc = None
        self._key_id = None
        self._key = None

    def set_key(self, key):
        self._key = RSA2048Key(key)
        self._key_id = SHA1(self._key)[-8:]

    def created(self):
        return self._key_id is not None

    @property
    def dc(self):
        return self._dc

    @dc.setter
    def dc(self, dc):
        self._dc = dc

    @property
    def key_id(self):
        if not self.created():
            raise MTProtoKeyNotReadyError("key_id")

        return self._key_id

    def prepare_aes(self, msg_key, send=True):
        """
        Defining AES Key and Initialization Vector

        The 2048-bit authorization key (auth_key) and the 128-bit message key (msg_key) are used to compute a
        56-bit AES key (aes_key) and a 256-bit initialization vector (aes_iv) which are subsequently used to
        encrypt the part of the message to be encrypted (i. e. everything with the exception of the external
            header which is added later) with AES-256 in infinite garble extension (IGE) mode.

        The algorithm for computing aes_key and aes_iv from auth_key and msg_key is as follows:

            sha1_a = SHA1 (msg_key + substr (auth_key, x, 32));
            sha1_b = SHA1 (substr (auth_key, 32+x, 16) + msg_key + substr (auth_key, 48+x, 16));
            sha1_с = SHA1 (substr (auth_key, 64+x, 32) + msg_key);
            sha1_d = SHA1 (msg_key + substr (auth_key, 96+x, 32));
            aes_key = substr (sha1_a, 0, 8) + substr (sha1_b, 8, 12) + substr (sha1_c, 4, 12);
            aes_iv = substr (sha1_a, 8, 12) + substr (sha1_b, 0, 8) + substr (sha1_c, 16, 4) + substr (sha1_d, 0, 8);

        where x = 0 for messages from client to server and x = 8 for those from server to client.

        The lower-order 1024 bits of auth_key are not involved in the computation. They may (together with the
        remaining bits or separately) be used on the client device to encrypt the local copy of the data
        received from the server. The 512 lower-order bits of auth_key are not stored on the server;
        therefore, if the client device uses them to encrypt local data and the user loses the key or the
        password, data decryption of local data is impossible (even if data from the server could be obtained).

        When AES is used to encrypt a block of data of a length not divisible by 16 bytes, the data is padded
        with random bytes to the smallest length divisible by 16 bytes immediately prior to being encrypted.
        """

        auth_key = self._key

        # TelegramDesktop uses condition "x = 8 if send else 0"
        # But that seems in contradiction to the statement above
        x = 0

        sha1_a = SHA1(msg_key + substr(auth_key, x, 32))
        sha1_b = SHA1(substr(auth_key, 32+x, 16) + msg_key + substr(auth_key, 48+x, 16))
        sha1_с = SHA1(substr(auth_key, 64+x, 32) + msg_key)
        sha1_d = SHA1(msg_key + substr(auth_key, 96+x, 32))
        aes_key = substr(sha1_a, 0, 8) + substr(sha1_b, 8, 12) + substr(sha1_с, 4, 12)
        aes_iv = substr(sha1_a, 8, 12) + substr(sha1_b, 0, 8) + substr(sha1_с, 16, 4) + substr(sha1_d, 0, 8)

        return AES256Key(aes_key), AES256Key(aes_iv)
<|endoftext|>


<commit_before>from __future__ import absolute_import

import logging

from keras.layers import Lambda, Dense, Concatenate
from keras.models import Model, Input

from .architectures import get_network_by_name

logger = logging.getLogger(__name__)

class BaseDecoder(object):
    def __init__(self, latent_dim, data_dim, network_architecture='synthetic', name='decoder'):
        logger.info("Initialising {} model with {}-dimensional data "
                    "and {}-dimensional latent input.".format(name, data_dim, latent_dim))
        self.name = '_'.join(name.lower().split())
        self.data_dim = data_dim
        self.latent_dim = latent_dim
        self.network_architecture = network_architecture
        self.data_input = Input(shape=(self.data_dim,), name='dec_ll_estimator_data_input')
        self.latent_input = Input(shape=(self.latent_dim,), name='dec_latent_input')

    def __call__(self, *args, **kwargs):
        return None

class StandardDecoder(BaseDecoder):
    """
    A StandardDecoder model has inputs comprising of a latent encoding given by an Encoder model, a prior sampler
    or other custom input and the raw Encoder data input, which is needed to estimate the reconstructed
    data log likelihood. It can be visualised as:

      Data    Latent
       |        |
       |    -----------
       |    | Decoder |
       |    -----------
       |        |
       |      Output
       |    probability    --->  Generated data
       |        |
       ---> Log Likelihood ---> -(reconstruction loss)

    Note that the reconstruction loss is not used when the model training ends. It serves only the purpose to
    define a measure of loss which is optimised.
    """
    def __init__(self, latent_dim, data_dim, network_architecture='synthetic', name=None):
        """
        Args:
            latent_dim: int, the flattened dimensionality of the latent space
            data_dim: int, the flattened dimensionality of the output space (data space)
            network_architecture: str, the architecture name for the body of the StandardDecoder model
            name: str, identifier of the model
        """
        super(StandardDecoder, self).__init__(latent_dim=latent_dim, data_dim=data_dim,
                                              network_architecture=network_architecture,
                                              name=name or 'Standard Decoder')

        generator_body = get_network_by_name['decoder'][network_architecture](self.latent_input, name_prefix=self.name)

        # NOTE: all decoder layers have names prefixed by `dec`.
        # This is essential for the partial model freezing during training.
        sampler_params = Dense(self.data_dim, activation='sigmoid',
                               name=self.name + 'dec_sampler_params')(generator_body)

        # a probability clipping is necessary for the Bernoulli `log_prob` property produces NaNs in the border cases.
        sampler_params = Lambda(lambda x: 1e-6 + (1 - 2e-6) * x,
                                name=self.name + 'dec_probs_clipper')(sampler_params)

        def bernoulli_log_probs(args):
            from tensorflow.contrib.distributions import Bernoulli
            mu, x = args
            log_px = Bernoulli(probs=mu, name=self.name + 'dec_bernoulli').log_prob(x)
            return log_px

        log_probs = Lambda(bernoulli_log_probs, name=self.name + 'dec_bernoulli_logprob')([sampler_params,
                                                                                self.data_input])

        self.generator = Model(inputs=self.latent_input, outputs=sampler_params, name=self.name + 'dec_sampling')
        self.ll_estimator = Model(inputs=[self.data_input, self.latent_input], outputs=log_probs,
                                  name=self.name + 'dec_trainable')

    def __call__(self, *args, **kwargs):
        """
        Make the StandardDecoder model callable on lists of Input layers or tensors.

        Args:
            *args: a list of input layers or tensors or numpy arrays, or a single input layer, tensor or numpy array.
        Keyword Args:
            is_learning: bool, whether the model is used for training or data generation. The output is either
                the reconstruction log likelihood or the output probabilities in the data space respectively.

        Returns:
            A StandardDecoder model in `training` or `data generation` mode.
        """
        is_learninig = kwargs.get('is_learning', True)
        if is_learninig:
            return self.ll_estimator(args[0])
        else:
            return self.generator(args[0])

class ConjointDecoder(object):
    """
    A ConjointDecoder model takes multiple latent dimensions and from each generates a data output. There is also a
    shared latent input which is used by all decoders:

               Latent_1   Latent_shared   Latent_1
                  |       /           \      |
             -------------             -------------
             | Decoder_1 |             | Decoder_2 |
             -------------             -------------
                  |                         |
               Output_1                  Output_2

    For more details about the loss computation, see the StandardDecoder class.
    """

    def __init__(self, latent_dims, data_dims, network_architecture='synthetic'):
        """
        Args:
            latent_dims: tuple, the flattened dimensions of the private latent spaces and the shared latent space
            data_dims: int, the flattened dimensions of the output spaces (data spaces)
            network_architecture: str, the architecture name for the body of all the decoders
        """
        assert len(latent_dims) == len(data_dims) + 1, \
            "Expected too receive {} private latent spaces and one shared for {} data inputs " \
            "but got {} instead.".format(len(data_dims) + 1, len(data_dims), len(latent_dims))

        name = "Conjoint Decoder"
        logger.info("Initialising {} model with {}-dimensional outputs "
                    "and {}-dimensional latent inputs.".format(name, data_dims, latent_dims))

        def bernoulli_log_probs(args):
            from tensorflow.contrib.distributions import Bernoulli
            mu, x = args
            log_px = Bernoulli(probs=mu, name='dec_bernoulli').log_prob(x)
            return log_px

        data_inputs = [Input(shape=(d,), name='dec_data_input_{}'.format(i)) for i, d in enumerate(data_dims)]
        latent_input = Input(shape=(sum(latent_dims),), name='dec_latent_input')

        # evt. set the output_shape arg.
        def slice_latent(x, start=0, stop=None):
            if stop is None:
                return x[:, start:]
            else:
                return x[:, start:stop]
        shared_latent_factors = Lambda(slice_latent, arguments={'start': -latent_dims[-1], 'stop': None})(latent_input)

        sampling_outputs, log_probs = [], []
        stop_id = 0
        for i in range(len(data_dims)):
            start_id = stop_id
            stop_id += latent_dims[i]
            latent_i = Lambda(slice_latent, arguments={'start': start_id, 'stop': stop_id})(latent_input)
            latent_i = Concatenate(axis=-1, name='dec_merged_latent_{}'.format(i))([latent_i, shared_latent_factors])
            generator_body = get_network_by_name['conjoint_decoder'][network_architecture](latent_i,
                                                                                'conj_{}'.format(i))
            sampler_params = Dense(data_dims[i], activation='sigmoid',
                                   name='dec_sampler_params_{}'.format(i))(generator_body)
            # probability clipping is necessary for the Bernoulli `log_prob` property produces NaNs in the limit cases
            sampler_params = Lambda(lambda x: 1e-6 + (1 - 2e-6) * x,
                                    name='dec_probs_clipper_{}'.format(i))(sampler_params)
            log_prob = Lambda(bernoulli_log_probs,
                              name='dec_bernoulli_logprob_{}'.format(i))([sampler_params, data_inputs[i]])
            sampling_outputs.append(sampler_params)
            log_probs.append(log_prob)

        log_probs = Concatenate(axis=-1, name='dec_conc_log_probs')(log_probs)

        self.generator = Model(inputs=latent_input, outputs=sampling_outputs, name='dec_sampling')
        self.ll_estimator = Model(inputs=data_inputs + [latent_input], outputs=log_probs, name='dec_trainable')

    def __call__(self, *args, **kwargs):
        """
        Make the StandardDecoder model callable on lists of Input layers or tensors.

        Args:
            *args: a list of input layers or tensors or numpy arrays, or a single input layer, tensor or numpy array.
        Keyword Args:
            is_learning: bool, whether the model is used for training or data generation. The output is either
                the reconstruction log likelihood or the output probabilities in the data space respectively.

        Returns:
            A StandardDecoder model in `training` or `data generation` mode.
        """
        is_learninig = kwargs.get('is_learning', True)
        if is_learninig:
            return self.ll_estimator(args[0])
        else:
            return self.generator(args[0])
<commit_msg>minors<commit_after>from __future__ import absolute_import

import logging

from keras.layers import Lambda, Dense, Concatenate
from keras.models import Model, Input

from .architectures import get_network_by_name

logger = logging.getLogger(__name__)

class BaseDecoder(object):
    def __init__(self, latent_dim, data_dim, network_architecture='synthetic', name='decoder'):
        logger.info("Initialising {} model with {}-dimensional data "
                    "and {}-dimensional latent input.".format(name, data_dim, latent_dim))
        self.name = '_'.join(name.lower().split())
        self.data_dim = data_dim
        self.latent_dim = latent_dim
        self.network_architecture = network_architecture
        self.data_input = Input(shape=(self.data_dim,), name='dec_ll_estimator_data_input')
        self.latent_input = Input(shape=(self.latent_dim,), name='dec_latent_input')

    def __call__(self, *args, **kwargs):
        return None

class StandardDecoder(BaseDecoder):
    """
    A StandardDecoder model has inputs comprising of a latent encoding given by an Encoder model, a prior sampler
    or other custom input and the raw Encoder data input, which is needed to estimate the reconstructed
    data log likelihood. It can be visualised as:

      Data    Latent
       |        |
       |    -----------
       |    | Decoder |
       |    -----------
       |        |
       |      Output
       |    probability    --->  Generated data
       |        |
       ---> Log Likelihood ---> -(reconstruction loss)

    Note that the reconstruction loss is not used when the model training ends. It serves only the purpose to
    define a measure of loss which is optimised.
    """
    def __init__(self, latent_dim, data_dim, network_architecture='synthetic', name=None):
        """
        Args:
            latent_dim: int, the flattened dimensionality of the latent space
            data_dim: int, the flattened dimensionality of the output space (data space)
            network_architecture: str, the architecture name for the body of the StandardDecoder model
            name: str, identifier of the model
        """
        super(StandardDecoder, self).__init__(latent_dim=latent_dim, data_dim=data_dim,
                                              network_architecture=network_architecture,
                                              name=name or 'Standard Decoder')

        generator_body = get_network_by_name['decoder'][network_architecture](self.latent_input, name_prefix=self.name)

        # NOTE: all decoder layers have names prefixed by `dec`.
        # This is essential for the partial model freezing during training.
        sampler_params = Dense(self.data_dim, activation='sigmoid',
                               name=self.name + 'dec_sampler_params')(generator_body)

        # a probability clipping is necessary for the Bernoulli `log_prob` property produces NaNs in the border cases.
        sampler_params = Lambda(lambda x: 1e-6 + (1 - 2e-6) * x,
                                name=self.name + 'dec_probs_clipper')(sampler_params)

        def bernoulli_log_probs(args):
            from tensorflow.contrib.distributions import Bernoulli
            mu, x = args
            log_px = Bernoulli(probs=mu, name=self.name + 'dec_bernoulli').log_prob(x)
            return log_px

        log_probs = Lambda(bernoulli_log_probs, name=self.name + 'dec_bernoulli_logprob')([sampler_params,
                                                                                self.data_input])

        self.generator = Model(inputs=self.latent_input, outputs=sampler_params, name=self.name + 'dec_sampling')
        self.ll_estimator = Model(inputs=[self.data_input, self.latent_input], outputs=log_probs,
                                  name=self.name + 'dec_trainable')

    def __call__(self, *args, **kwargs):
        """
        Make the StandardDecoder model callable on lists of Input layers or tensors.

        Args:
            *args: a list of input layers or tensors or numpy arrays, or a single input layer, tensor or numpy array.
        Keyword Args:
            is_learning: bool, whether the model is used for training or data generation. The output is either
                the reconstruction log likelihood or the output probabilities in the data space respectively.

        Returns:
            A StandardDecoder model in `training` or `data generation` mode.
        """
        is_learninig = kwargs.get('is_learning', True)
        if is_learninig:
            return self.ll_estimator(args[0])
        else:
            return self.generator(args[0])

class ConjointDecoder(object):
    """
    A ConjointDecoder model takes multiple latent dimensions and from each generates a data output. There is also a
    shared latent input which is used by all decoders:

               Latent_1   Latent_shared   Latent_1
                  |       /           \      |
             -------------             -------------
             | Decoder_1 |             | Decoder_2 |
             -------------             -------------
                  |                         |
               Output_1                  Output_2

    For more details about the loss computation, see the StandardDecoder class.
    """

    def __init__(self, latent_dims, data_dims, network_architecture='synthetic'):
        """
        Args:
            latent_dims: tuple, the flattened dimensions of the private latent spaces and the shared latent space
            data_dims: int, the flattened dimensions of the output spaces (data spaces)
            network_architecture: str, the architecture name for the body of all the decoders
        """
        assert len(latent_dims) == len(data_dims) + 1, \
            "Expected too receive {} private latent spaces and one shared for {} data inputs " \
            "but got {} instead.".format(len(data_dims) + 1, len(data_dims), len(latent_dims))

        name = "Conjoint Decoder"
        logger.info("Initialising {} model with {}-dimensional outputs "
                    "and {}-dimensional latent inputs.".format(name, data_dims, latent_dims))

        def bernoulli_log_probs(args):
            from tensorflow.contrib.distributions import Bernoulli
            mu, x = args
            log_px = Bernoulli(probs=mu, name='dec_bernoulli').log_prob(x)
            return log_px

        data_inputs = [Input(shape=(d,), name='dec_data_input_{}'.format(i)) for i, d in enumerate(data_dims)]
        latent_input = Input(shape=(sum(latent_dims),), name='dec_latent_input')

        # evt. set the output_shape arg.
        def slice_latent(x, start=0, stop=None):
            if stop is None:
                return x[:, start:]
            else:
                return x[:, start:stop]
        shared_latent_factors = Lambda(slice_latent, arguments={'start': -latent_dims[-1], 'stop': None},
                                       name='dec_slice_shared_lat')(latent_input)

        sampling_outputs, log_probs = [], []
        stop_id = 0
        for i in range(len(data_dims)):
            start_id = stop_id
            stop_id += latent_dims[i]
            latent_i = Lambda(slice_latent, arguments={'start': start_id, 'stop': stop_id},
                              name='dec_slice_{}'.format(i))(latent_input)
            latent_i = Concatenate(axis=-1, name='dec_merged_latent_{}'.format(i))([latent_i, shared_latent_factors])
            generator_body = get_network_by_name['conjoint_decoder'][network_architecture](latent_i,
                                                                                'conj_{}'.format(i))
            sampler_params = Dense(data_dims[i], activation='sigmoid',
                                   name='dec_sampler_params_{}'.format(i))(generator_body)
            # probability clipping is necessary for the Bernoulli `log_prob` property produces NaNs in the limit cases
            sampler_params = Lambda(lambda x: 1e-6 + (1 - 2e-6) * x,
                                    name='dec_probs_clipper_{}'.format(i))(sampler_params)
            log_prob = Lambda(bernoulli_log_probs,
                              name='dec_bernoulli_logprob_{}'.format(i))([sampler_params, data_inputs[i]])
            sampling_outputs.append(sampler_params)
            log_probs.append(log_prob)

        log_probs = Concatenate(axis=-1, name='dec_conc_log_probs')(log_probs)

        self.generator = Model(inputs=latent_input, outputs=sampling_outputs, name='dec_sampling')
        self.ll_estimator = Model(inputs=data_inputs + [latent_input], outputs=log_probs, name='dec_trainable')

    def __call__(self, *args, **kwargs):
        """
        Make the StandardDecoder model callable on lists of Input layers or tensors.

        Args:
            *args: a list of input layers or tensors or numpy arrays, or a single input layer, tensor or numpy array.
        Keyword Args:
            is_learning: bool, whether the model is used for training or data generation. The output is either
                the reconstruction log likelihood or the output probabilities in the data space respectively.

        Returns:
            A StandardDecoder model in `training` or `data generation` mode.
        """
        is_learninig = kwargs.get('is_learning', True)
        if is_learninig:
            return self.ll_estimator(args[0])
        else:
            return self.generator(args[0])
<|endoftext|>