<commit_before>"""
The MIT License (MIT)

Copyright (c) 2017 Stefan Graupner

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import sys
from threading import activeCount, Thread
from time import sleep

from .client import Client
from .entity_queue import EntityQueue
from .exceptions import \
    EndpointNotReachableException, \
    EndpointIsNotAnOParlEndpointException, \
    ObjectValidationFailedException
from .output import Output
from .result import Result
from .seen_list import SeenList

VALID_OPARL_VERSIONS = [
    "https://schema.oparl.org/1.0/"
]

class Validator:
    """
        This class provides a Validator instance.

        The OParl Validator always requires an endpoint to validate
        and an optional set of options which, if given, should
        be an argparse.Namespace object.

        Please refer to `self.parse_options` for a listing of
        currently available options.

        In order to run, the Validator requires a redis connection
        which currently must be on the same machine.
    """

    NUM_VALIDATION_WORKERS = 3

    def __init__(self, endpoint, options = None):
        self.endpoint = endpoint
        self.options = self.parse_options(options)

        Output.initialize(
            self.options.porcelain,
            self.options.silent,
            self.options.verbosity
        )

        if self.options.porcelain and self.options.result is None:
            Output.message('You selected an invalid option, please provide a result destination filename')
            exit(1)

    def parse_options(self, options):
        if 'format' not in options:
            options.format = 'json'

        if 'result' not in options:
            options.output = None

        if 'porcelain' not in options:
            options.porcelain = False

        if 'silent' not in options:
            options.silent = True

        if 'read' not in options:
            options.read = False

        return options

    def run(self):
        result = None

        if self.options.read:
            result = Result.from_file(self.endpoint)
        else:
            try:
                self.client = Client(self.endpoint)
            except EndpointNotReachableException:
                Output.message('Endpoint {} is not reachable, aborting validation.', self.endpoint)
                exit(1)
            except EndpointIsNotAnOParlEndpointException:
                Output.message('Endpoint {} is not an OParl-Endpoint, aborting validation.', self.endpoint)
                exit(1)

            result = self.validate()
            result.compile()

        self.handle_result(result)

    def validate(self):
        Output.message("Beginning validation of {}", self.endpoint)
        Output.message("Found '{}'", self.client.system.get_name())

        bodies = self.client.system.get_body()
        num_bodies = len(bodies)

        unprocessed_entities = EntityQueue(maxsize = 1000)

        seen_list = SeenList()
        result = Result()

        walker_threads = []
        worker_threads = []

        for body in bodies:
            walker = self.client.create_body_walker(body, unprocessed_entities)
            walker_threads.append(walker)

        for thread in walker_threads:
            thread.start()

        # let the walkers walk a little
        while unprocessed_entities.empty():
            pass

        # TODO: make this dependent on system resources / an option

        Output.add_progress_bar('validation_progress', 'Validating')

        for i in range(0, Validator.NUM_VALIDATION_WORKERS):
            worker = ValidationWorker(
                'validation_worker_{}'.format(i),
                unprocessed_entities,
                seen_list,
                result
            )
            worker_threads.append(worker)

        for thread in worker_threads:
            thread.start()

        for thread in walker_threads:
            thread.join()

        while not unprocessed_entities.empty():
            pass

        for thread in worker_threads:
            thread.join()

        Output.message("Validation finished")

        result.network = self.client.network
        result.total_entities = len(seen_list)

        return result

    def handle_result(self, result):
        formatted_result = result.text()
        if self.options.format == 'json':
            formatted_result = result.json()

        if Output.porcelain or self.options.result is not None:
            with open(self.options.result, 'w+') as f:
                f.write(formatted_result)
                exit(0)

        print(formatted_result)

class ValidationWorker(Thread):
    def __init__(self, id, queue, seen_list, result):
        super(ValidationWorker, self).__init__()
        # TODO: add progress reporters

        self.id = id
        self.queue = queue
        self.result = result
        self.seen_list = seen_list
        self.current_object = None

    def run(self):
        while not self.queue.empty():
            self.get_next_object()

            if not self.is_seen_object():
                try:
                    results = self.validate_object()
                    self.save_validation_results(results)
                except ObjectValidationFailedException:
                    self.save_failed_object()

                Output.update_progress_bar('validation_progress', remaining=self.queue.qsize())

    def get_next_object(self):
        self.queue.acquire()
        self.current_object = self.queue.get()
        self.queue.release()

    def is_seen_object(self):
        if self.current_object.get_id() in self.seen_list:
            return True

        self.seen_list.push(self.current_object.get_id())

        return False

    def validate_object(self):
        try:
            validation_results = self.current_object.validate()
        except GLib.Error as glib_error:
            # TODO: track liboparl validation errors
            raise ObjectValidationFailedException from glib_error

        # TODO: implement additional checks like e.g. file reachability

        return validation_results

    def save_validation_results(self, validation_results):
        self.result.acquire()

        if len(validation_results) > 0:
            self.result.failed_entities += 1

        for result in validation_results:
            self.result.parse_validation_result(self.current_object, result)

        self.result.release()

    def save_failed_object(self):
        self.result.acquire()
        self.result.fatal_objects.append(self.current_object)
        self.result.release()<commit_msg>Removed TODO<commit_after>"""
The MIT License (MIT)

Copyright (c) 2017 Stefan Graupner

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import sys
from threading import activeCount, Thread
from time import sleep

from .client import Client
from .entity_queue import EntityQueue
from .exceptions import \
    EndpointNotReachableException, \
    EndpointIsNotAnOParlEndpointException, \
    ObjectValidationFailedException
from .output import Output
from .result import Result
from .seen_list import SeenList

VALID_OPARL_VERSIONS = [
    "https://schema.oparl.org/1.0/"
]

class Validator:
    """
        This class provides a Validator instance.

        The OParl Validator always requires an endpoint to validate
        and an optional set of options which, if given, should
        be an argparse.Namespace object.

        Please refer to `self.parse_options` for a listing of
        currently available options.

        In order to run, the Validator requires a redis connection
        which currently must be on the same machine.
    """

    NUM_VALIDATION_WORKERS = 3

    def __init__(self, endpoint, options = None):
        self.endpoint = endpoint
        self.options = self.parse_options(options)

        Output.initialize(
            self.options.porcelain,
            self.options.silent,
            self.options.verbosity
        )

        if self.options.porcelain and self.options.result is None:
            Output.message('You selected an invalid option, please provide a result destination filename')
            exit(1)

    def parse_options(self, options):
        if 'format' not in options:
            options.format = 'json'

        if 'result' not in options:
            options.output = None

        if 'porcelain' not in options:
            options.porcelain = False

        if 'silent' not in options:
            options.silent = True

        if 'read' not in options:
            options.read = False

        return options

    def run(self):
        result = None

        if self.options.read:
            result = Result.from_file(self.endpoint)
        else:
            try:
                self.client = Client(self.endpoint)
            except EndpointNotReachableException:
                Output.message('Endpoint {} is not reachable, aborting validation.', self.endpoint)
                exit(1)
            except EndpointIsNotAnOParlEndpointException:
                Output.message('Endpoint {} is not an OParl-Endpoint, aborting validation.', self.endpoint)
                exit(1)

            result = self.validate()
            result.compile()

        self.handle_result(result)

    def validate(self):
        Output.message("Beginning validation of {}", self.endpoint)
        Output.message("Found '{}'", self.client.system.get_name())

        bodies = self.client.system.get_body()
        num_bodies = len(bodies)

        unprocessed_entities = EntityQueue(maxsize = 1000)

        seen_list = SeenList()
        result = Result()

        walker_threads = []
        worker_threads = []

        for body in bodies:
            walker = self.client.create_body_walker(body, unprocessed_entities)
            walker_threads.append(walker)

        for thread in walker_threads:
            thread.start()

        # let the walkers walk a little
        while unprocessed_entities.empty():
            pass

        # TODO: make this dependent on system resources / an option

        Output.add_progress_bar('validation_progress', 'Validating')

        for i in range(0, Validator.NUM_VALIDATION_WORKERS):
            worker = ValidationWorker(
                'validation_worker_{}'.format(i),
                unprocessed_entities,
                seen_list,
                result
            )
            worker_threads.append(worker)

        for thread in worker_threads:
            thread.start()

        for thread in walker_threads:
            thread.join()

        while not unprocessed_entities.empty():
            pass

        for thread in worker_threads:
            thread.join()

        Output.message("Validation finished")

        result.network = self.client.network
        result.total_entities = len(seen_list)

        return result

    def handle_result(self, result):
        formatted_result = result.text()
        if self.options.format == 'json':
            formatted_result = result.json()

        if Output.porcelain or self.options.result is not None:
            with open(self.options.result, 'w+') as f:
                f.write(formatted_result)
                exit(0)

        print(formatted_result)

class ValidationWorker(Thread):
    def __init__(self, id, queue, seen_list, result):
        super(ValidationWorker, self).__init__()
        self.id = id
        self.queue = queue
        self.result = result
        self.seen_list = seen_list
        self.current_object = None

    def run(self):
        while not self.queue.empty():
            self.get_next_object()

            if not self.is_seen_object():
                try:
                    results = self.validate_object()
                    self.save_validation_results(results)
                except ObjectValidationFailedException:
                    self.save_failed_object()

                Output.update_progress_bar('validation_progress', remaining=self.queue.qsize())

    def get_next_object(self):
        self.queue.acquire()
        self.current_object = self.queue.get()
        self.queue.release()

    def is_seen_object(self):
        if self.current_object.get_id() in self.seen_list:
            return True

        self.seen_list.push(self.current_object.get_id())

        return False

    def validate_object(self):
        try:
            validation_results = self.current_object.validate()
        except GLib.Error as glib_error:
            # TODO: track liboparl validation errors
            raise ObjectValidationFailedException from glib_error

        # TODO: implement additional checks like e.g. file reachability

        return validation_results

    def save_validation_results(self, validation_results):
        self.result.acquire()

        if len(validation_results) > 0:
            self.result.failed_entities += 1

        for result in validation_results:
            self.result.parse_validation_result(self.current_object, result)

        self.result.release()

    def save_failed_object(self):
        self.result.acquire()
        self.result.fatal_objects.append(self.current_object)
        self.result.release()<|endoftext|>


<commit_before>import json
from crdt import CRDT
from collections import Counter

"""
operation       | item in external set | item not in external set
----------------+----------------------+--------------------------
('a', 'd', 'd') | delete item          | delete item (no op)
('d', 'a', 'd') | delete item          | delete item (no op)
('d', 'd', 'a') | add item             | add item
('a', 'a', 'd') | delete item          | delete item (no op)
('a', 'd', 'a') | no op                | add item
('d', 'a', 'a') | no op                | add item
"""

class DistributedSet(CRDT):
    def __init__(self, name, url, db, auth=(), params={}, always_write=False):
        super(DistributedSet, self).__init__(
            name,
            url,
            db,
            auth,
            params,
            always_write
        )
        if 'cloudant.com' not in url:
            del self.params['key']
            self.params['startkey'] = json.dumps([name])
            self.params['endkey'] = json.dumps([name, {}])
        self.additions = Counter()
        self.deletions = Counter()
        self.value = set()
        self.default_state = set()

    def add(self, item):
        self.additions[item] += 1

    def remove(self, item):
        self.deletions[item] += 1

    def __iter__(self):
        """
        Do the set operation and return the iterable over the result
        """
        adds, deletes = self._get_state()
        return iter(
            self.value.union(set(adds)) - set(deletes)
        )

    def __str__(self):
        adds, deletes = self._get_state()
        return str(
            self.value.union(set(adds)) - set(deletes)
        )

    def _update(self, item):
        """
        The _update method on a DistributedSet is only used to reset internal
        state to their original state.
        """
        self.additions = Counter()
        self.deletions = Counter()

    def _get_state(self):
        adds = self.additions - self.deletions
        deletes = self.deletions - self.additions
        return adds, Counter(dict((k, -1 * v) for k, v in deletes.iteritems()))

    def _serialise(self):
        adds, deletes = self._get_state()
        return {
            'additions': adds,
            'deletions': deletes
        }

    def _parse(self, data):
        if 'cloudant.com' in data.url:
            data = data.json()['rows'][0]['value']
            return set([k for k, v in data.iteritems() if v > 0])
        else:
            s = set()
            for row in data.json()['rows']:
                if row['value'] > 0:
                    s.add(row['key'][1])
            return s
<commit_msg>comment on why adding up deletions<commit_after>import json
from crdt import CRDT
from collections import Counter

"""
operation       | item in external set | item not in external set
----------------+----------------------+--------------------------
('a', 'd', 'd') | delete item          | delete item (no op)
('d', 'a', 'd') | delete item          | delete item (no op)
('d', 'd', 'a') | add item             | add item
('a', 'a', 'd') | delete item          | delete item (no op)
('a', 'd', 'a') | no op                | add item
('d', 'a', 'a') | no op                | add item
"""

class DistributedSet(CRDT):
    def __init__(self, name, url, db, auth=(), params={}, always_write=False):
        super(DistributedSet, self).__init__(
            name,
            url,
            db,
            auth,
            params,
            always_write
        )
        if 'cloudant.com' not in url:
            del self.params['key']
            self.params['startkey'] = json.dumps([name])
            self.params['endkey'] = json.dumps([name, {}])
        self.additions = Counter()
        self.deletions = Counter()
        self.value = set()
        self.default_state = set()

    def add(self, item):
        self.additions[item] += 1

    def remove(self, item):
        # Have to add here because we use counter subtraction below
        self.deletions[item] += 1

    def __iter__(self):
        """
        Do the set operation and return the iterable over the result
        """
        adds, deletes = self._get_state()
        return iter(
            self.value.union(set(adds)) - set(deletes)
        )

    def __str__(self):
        adds, deletes = self._get_state()
        return str(
            self.value.union(set(adds)) - set(deletes)
        )

    def _update(self, item):
        """
        The _update method on a DistributedSet is only used to reset internal
        state to their original state.
        """
        self.additions = Counter()
        self.deletions = Counter()

    def _get_state(self):
        adds = self.additions - self.deletions
        deletes = self.deletions - self.additions
        return adds, Counter(dict((k, -1 * v) for k, v in deletes.iteritems()))

    def _serialise(self):
        adds, deletes = self._get_state()
        return {
            'additions': adds,
            'deletions': deletes
        }

    def _parse(self, data):
        if 'cloudant.com' in data.url:
            data = data.json()['rows'][0]['value']
            return set([k for k, v in data.iteritems() if v > 0])
        else:
            s = set()
            for row in data.json()['rows']:
                if row['value'] > 0:
                    s.add(row['key'][1])
            return s
<|endoftext|>


<commit_before>from django import template
import json
import ntpath

register = template.Library()

@register.filter(name='n3')
def n3(filename):
    """
    :param filename:
    :return: file name with suffix .n3
    """

    return filename.rsplit(".", 1)[0] + ".n3"

def path_leaf(path):
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)

@register.filter(name='kill_path')
def kill_path(path):
    """
    :param path:
    :return: file without path
    """
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)

@register.filter(name='model_as_table')
def model_as_table(model, num_rows=-1):
    result = '
'
    result += model_header_as_table(model)
    result += model_content_as_table(model, num_rows)
    result += '
'
    return result

@register.filter(name='model_as_table2')
def model_as_table2(model, num_rows=-1):
    result = '
'
    result += model_header_as_table(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_table_predicate')
def model_as_table_predicate(model, num_rows=-1):
    result = '
'
    result += model_header_as_table_predicate(model)
    result += model_content_as_table(model, num_rows)
    result += '
'
    return result

@register.filter(name='model_as_table_predicate2')
def model_as_table_predicate2(model):
    result = '
'
    result += model_header_as_table_predicate(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_table_object')
def model_as_table_object(model, pagination):
    result = '
'
    result += model_header_as_table_object(model)
    result += model_content_as_table(model, pagination)
    result += '
'
    return result

@register.filter(name='model_as_table_object2')
def model_as_table_object2(model):
    result = '
'
    result += model_header_as_table_object(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_thead')
def model_header_as_table(model):
    model = json.loads(model)
    headers = []
    for col in model['columns']:
        if 'col_num_new' not in col or col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"
    result += "
"

    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"

    result += "
"
    result += "
"

    return result

@register.filter(name='model_as_tbody2')
def model_content_as_table2(model):

    model = json.loads(model)
    f = model['excerpt']['start_row']

    selected_cols = []

    for col in model['columns']:
        if col['col_num_new'] > -1:
            selected_cols.append(True)
        else:
            selected_cols.append(False)

    result = "
"
    for i, row in enumerate(model['excerpt']['rows']):
        result += "
"
        for j, field in enumerate(row):
            if selected_cols[j] is True:
                result += ''
                result += '' + field + ''
                result += "	"
        result += "
"
    result += "
"

    return result

@register.filter(name='model_as_tbody')
def model_content_as_table(model, pagination):
    '''
    pagination can be
    int number: first x elements will be selected
    dict of 'pagination', with page and perPage attributes as in views.py -> csv_object function
    '''
    model = json.loads(model)
    #num_rows = model['num_cols_selected']
    content = []
    #p = False
    f = 0

    if isinstance(pagination, dict) and 'page' in pagination and 'perPage' in pagination:
        #p = True
        f = (pagination['page'] - 1) * pagination['perPage']
        #t = f + pagination['perPage']

    for col in model['columns']:
        if 'col_num_new' not in col or col['col_num_new'] > -1:  # show column
            row = []

            #if p:
            #    fields = col['fields'][f:t]
            if isinstance(pagination, int):
                fields = col['fields'][:pagination]
            else:
                fields = col['fields']

            for elem in fields:
                row.append(elem['orig_val'])
            content.append(row)

    content = list(zip(*content))

    result = "
"

    for i, row in enumerate(content):
        result += "
"
        for j, field in enumerate(row):
            result += ''
            result += '' + field + ''
            result += "	"
        result += "
"

    result += "
"

    return result

@register.filter(name='model_as_thead_predicate')
def model_header_as_table_predicate(model):
    model = json.loads(model)
    # num_rows = model['num_cols_selected']
    headers = []
    for col in model['columns']:
        if not 'col_num_new' in col or col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"
    result += "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += '<input type="text" value="' + field + '" id="search_textinput_' + str(
            i + 1) + '" onkeyup="delayedOracleCall(' + str(i + 1) + ')" autocomplete="off">'
        result += ''
        result += '

'
        result += '
'
        result += '

'
        result += '	'
    result += "
"

    result += "
"

    return result

@register.filter(name='model_as_thead_object')
def model_header_as_table_object(model):
    model = json.loads(model)
    # num_rows = model['num_cols_selected']
    headers = []
    for col in model['columns']:
        if col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += '<select>'
        result += '<option>no action</option>'
        result += '<option>add URIs</option>'
        result += '<option>add data type</option>'
        result += '</select>'
        result += '	'
    result += "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"
    result += "
"

    result += "
"

    return result
<commit_msg>col count bug<commit_after>from django import template
import json
import ntpath

register = template.Library()

@register.filter(name='n3')
def n3(filename):
    """
    :param filename:
    :return: file name with suffix .n3
    """

    return filename.rsplit(".", 1)[0] + ".n3"

def path_leaf(path):
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)

@register.filter(name='kill_path')
def kill_path(path):
    """
    :param path:
    :return: file without path
    """
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)

@register.filter(name='model_as_table')
def model_as_table(model, num_rows=-1):
    result = '
'
    result += model_header_as_table(model)
    result += model_content_as_table(model, num_rows)
    result += '
'
    return result

@register.filter(name='model_as_table2')
def model_as_table2(model, num_rows=-1):
    result = '
'
    result += model_header_as_table(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_table_predicate')
def model_as_table_predicate(model, num_rows=-1):
    result = '
'
    result += model_header_as_table_predicate(model)
    result += model_content_as_table(model, num_rows)
    result += '
'
    return result

@register.filter(name='model_as_table_predicate2')
def model_as_table_predicate2(model):
    result = '
'
    result += model_header_as_table_predicate(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_table_object')
def model_as_table_object(model, pagination):
    result = '
'
    result += model_header_as_table_object(model)
    result += model_content_as_table(model, pagination)
    result += '
'
    return result

@register.filter(name='model_as_table_object2')
def model_as_table_object2(model):
    result = '
'
    result += model_header_as_table_object(model)
    result += model_content_as_table2(model)
    result += '
'
    return result

@register.filter(name='model_as_thead')
def model_header_as_table(model):
    model = json.loads(model)
    headers = []
    for col in model['columns']:
        if 'col_num_new' not in col or col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"
    result += "
"

    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"

    result += "
"
    result += "
"

    return result

@register.filter(name='model_as_tbody2')
def model_content_as_table2(model):

    model = json.loads(model)
    f = model['excerpt']['start_row']

    selected_cols = []

    for col in model['columns']:
        if col['col_num_new'] > -1:
            selected_cols.append(True)
        else:
            selected_cols.append(False)

    result = "
"
    for i, row in enumerate(model['excerpt']['rows']):
        result += "
"
        col_counter = 0
        for j, field in enumerate(row):
            if selected_cols[j] is True:
                col_counter += 1
                result += ''
                result += '' + field + ''
                result += "	"
        result += "
"
    result += "
"

    return result

@register.filter(name='model_as_tbody')
def model_content_as_table(model, pagination):
    '''
    pagination can be
    int number: first x elements will be selected
    dict of 'pagination', with page and perPage attributes as in views.py -> csv_object function
    '''
    model = json.loads(model)
    #num_rows = model['num_cols_selected']
    content = []
    #p = False
    f = 0

    if isinstance(pagination, dict) and 'page' in pagination and 'perPage' in pagination:
        #p = True
        f = (pagination['page'] - 1) * pagination['perPage']
        #t = f + pagination['perPage']

    for col in model['columns']:
        if 'col_num_new' not in col or col['col_num_new'] > -1:  # show column
            row = []

            #if p:
            #    fields = col['fields'][f:t]
            if isinstance(pagination, int):
                fields = col['fields'][:pagination]
            else:
                fields = col['fields']

            for elem in fields:
                row.append(elem['orig_val'])
            content.append(row)

    content = list(zip(*content))

    result = "
"

    for i, row in enumerate(content):
        result += "
"
        for j, field in enumerate(row):
            result += ''
            result += '' + field + ''
            result += "	"
        result += "
"

    result += "
"

    return result

@register.filter(name='model_as_thead_predicate')
def model_header_as_table_predicate(model):
    model = json.loads(model)
    # num_rows = model['num_cols_selected']
    headers = []
    for col in model['columns']:
        if not 'col_num_new' in col or col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"
    result += "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += '<input type="text" value="' + field + '" id="search_textinput_' + str(
            i + 1) + '" onkeyup="delayedOracleCall(' + str(i + 1) + ')" autocomplete="off">'
        result += ''
        result += '

'
        result += '
'
        result += '

'
        result += '	'
    result += "
"

    result += "
"

    return result

@register.filter(name='model_as_thead_object')
def model_header_as_table_object(model):
    model = json.loads(model)
    # num_rows = model['num_cols_selected']
    headers = []
    for col in model['columns']:
        if col['col_num_new'] > -1:  # show column
            headers.append(col['header']['orig_val'])

    result = "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += '<select>'
        result += '<option>no action</option>'
        result += '<option>add URIs</option>'
        result += '<option>add data type</option>'
        result += '</select>'
        result += '	'
    result += "
"

    result += "
"
    for i, field in enumerate(headers):
        result += ''
        result += field
        result += "	"
    result += "
"

    result += "
"

    return result
<|endoftext|>


<commit_before>import time
import random
import datetime
from unittest import mock

import pytest
import pytz
import freezegun

from tempora import schedule

do_nothing = type(None)

def test_delayed_command_order():
    """
    delayed commands should be sorted by delay time
    """
    delays = [random.randint(0, 99) for x in range(5)]
    cmds = sorted(
        [schedule.DelayedCommand.after(delay, do_nothing) for delay in delays]
    )
    assert [c.delay.seconds for c in cmds] == sorted(delays)

def test_periodic_command_delay():
    "A PeriodicCommand must have a positive, non-zero delay."
    with pytest.raises(ValueError) as exc_info:
        schedule.PeriodicCommand.after(0, None)
    assert str(exc_info.value) == test_periodic_command_delay.__doc__

def test_periodic_command_fixed_delay():
    """
    Test that we can construct a periodic command with a fixed initial
    delay.
    """
    fd = schedule.PeriodicCommandFixedDelay.at_time(
        at=schedule.now(), delay=datetime.timedelta(seconds=2), target=lambda: None
    )
    assert fd.due() is True
    assert fd.next().due() is False

class TestCommands:
    def test_delayed_command_from_timestamp(self):
        """
        Ensure a delayed command can be constructed from a timestamp.
        """
        t = time.time()
        schedule.DelayedCommand.at_time(t, do_nothing)

    def test_command_at_noon(self):
        """
        Create a periodic command that's run at noon every day.
        """
        when = datetime.time(12, 0, tzinfo=pytz.utc)
        cmd = schedule.PeriodicCommandFixedDelay.daily_at(when, target=None)
        assert cmd.due() is False
        next_cmd = cmd.next()
        daily = datetime.timedelta(days=1)
        day_from_now = schedule.now() + daily
        two_days_from_now = day_from_now + daily
        assert day_from_now < next_cmd < two_days_from_now

    @pytest.mark.parametrize("hour", range(24))
    def test_command_at_noon_distant_local(self, hour):
        """
        Run test_command_at_noon, but with the local timezone
        more than 12 hours away from UTC.
        """
        with freezegun.freeze_time(f"2020-01-10 {hour:02}:01", tz_offset=-14):
            self.test_command_at_noon()

class TestTimezones:
    def test_alternate_timezone_west(self):
        target_tz = pytz.timezone('US/Pacific')
        target = schedule.now().astimezone(target_tz)
        cmd = schedule.DelayedCommand.at_time(target, target=None)
        assert cmd.due()

    def test_alternate_timezone_east(self):
        target_tz = pytz.timezone('Europe/Amsterdam')
        target = schedule.now().astimezone(target_tz)
        cmd = schedule.DelayedCommand.at_time(target, target=None)
        assert cmd.due()

    def test_daylight_savings(self):
        """
        A command at 9am should always be 9am regardless of
        a DST boundary.
        """
        with freezegun.freeze_time('2018-03-10 08:00:00'):
            target_tz = pytz.timezone('US/Eastern')
            target_time = datetime.time(9, tzinfo=target_tz)
            cmd = schedule.PeriodicCommandFixedDelay.daily_at(
                target_time, target=lambda: None
            )

        def naive(dt):
            return dt.replace(tzinfo=None)

        assert naive(cmd) == datetime.datetime(2018, 3, 10, 9, 0, 0)
        next_ = cmd.next()
        assert naive(next_) == datetime.datetime(2018, 3, 11, 9, 0, 0)
        assert next_ - cmd == datetime.timedelta(hours=23)

class TestScheduler:
    def test_invoke_scheduler(self):
        sched = schedule.InvokeScheduler()
        target = mock.MagicMock()
        cmd = schedule.DelayedCommand.after(0, target)
        sched.add(cmd)
        sched.run_pending()
        target.assert_called_once()
        assert not sched.queue

    def test_callback_scheduler(self):
        callback = mock.MagicMock()
        sched = schedule.CallbackScheduler(callback)
        target = mock.MagicMock()
        cmd = schedule.DelayedCommand.after(0, target)
        sched.add(cmd)
        sched.run_pending()
        callback.assert_called_once_with(target)

    def test_periodic_command(self):
        sched = schedule.InvokeScheduler()
        target = mock.MagicMock()
        cmd = schedule.PeriodicCommand.at_time(time.time() + 0.1, target)
        sched.add(cmd)
        sched.run_pending()
        target.assert_not_called()
        time.sleep(0.1)
        sched.run_pending()
        assert sched.queue
        target.assert_called_once()
        time.sleep(0.1)
        sched.run_pending()
        assert target.call_count == 2
<commit_msg>Test both +14 and -14 timezone offsets. Ref #8.<commit_after>import time
import random
import datetime
from unittest import mock

import pytest
import pytz
import freezegun

from tempora import schedule

do_nothing = type(None)

def test_delayed_command_order():
    """
    delayed commands should be sorted by delay time
    """
    delays = [random.randint(0, 99) for x in range(5)]
    cmds = sorted(
        [schedule.DelayedCommand.after(delay, do_nothing) for delay in delays]
    )
    assert [c.delay.seconds for c in cmds] == sorted(delays)

def test_periodic_command_delay():
    "A PeriodicCommand must have a positive, non-zero delay."
    with pytest.raises(ValueError) as exc_info:
        schedule.PeriodicCommand.after(0, None)
    assert str(exc_info.value) == test_periodic_command_delay.__doc__

def test_periodic_command_fixed_delay():
    """
    Test that we can construct a periodic command with a fixed initial
    delay.
    """
    fd = schedule.PeriodicCommandFixedDelay.at_time(
        at=schedule.now(), delay=datetime.timedelta(seconds=2), target=lambda: None
    )
    assert fd.due() is True
    assert fd.next().due() is False

class TestCommands:
    def test_delayed_command_from_timestamp(self):
        """
        Ensure a delayed command can be constructed from a timestamp.
        """
        t = time.time()
        schedule.DelayedCommand.at_time(t, do_nothing)

    def test_command_at_noon(self):
        """
        Create a periodic command that's run at noon every day.
        """
        when = datetime.time(12, 0, tzinfo=pytz.utc)
        cmd = schedule.PeriodicCommandFixedDelay.daily_at(when, target=None)
        assert cmd.due() is False
        next_cmd = cmd.next()
        daily = datetime.timedelta(days=1)
        day_from_now = schedule.now() + daily
        two_days_from_now = day_from_now + daily
        assert day_from_now < next_cmd < two_days_from_now

    @pytest.mark.parametrize("hour", range(10, 14))
    @pytest.mark.parametrize("tz_offset", (14, -14))
    def test_command_at_noon_distant_local(self, hour, tz_offset):
        """
        Run test_command_at_noon, but with the local timezone
        more than 12 hours away from UTC.
        """
        with freezegun.freeze_time(f"2020-01-10 {hour:02}:01", tz_offset=tz_offset):
            self.test_command_at_noon()

class TestTimezones:
    def test_alternate_timezone_west(self):
        target_tz = pytz.timezone('US/Pacific')
        target = schedule.now().astimezone(target_tz)
        cmd = schedule.DelayedCommand.at_time(target, target=None)
        assert cmd.due()

    def test_alternate_timezone_east(self):
        target_tz = pytz.timezone('Europe/Amsterdam')
        target = schedule.now().astimezone(target_tz)
        cmd = schedule.DelayedCommand.at_time(target, target=None)
        assert cmd.due()

    def test_daylight_savings(self):
        """
        A command at 9am should always be 9am regardless of
        a DST boundary.
        """
        with freezegun.freeze_time('2018-03-10 08:00:00'):
            target_tz = pytz.timezone('US/Eastern')
            target_time = datetime.time(9, tzinfo=target_tz)
            cmd = schedule.PeriodicCommandFixedDelay.daily_at(
                target_time, target=lambda: None
            )

        def naive(dt):
            return dt.replace(tzinfo=None)

        assert naive(cmd) == datetime.datetime(2018, 3, 10, 9, 0, 0)
        next_ = cmd.next()
        assert naive(next_) == datetime.datetime(2018, 3, 11, 9, 0, 0)
        assert next_ - cmd == datetime.timedelta(hours=23)

class TestScheduler:
    def test_invoke_scheduler(self):
        sched = schedule.InvokeScheduler()
        target = mock.MagicMock()
        cmd = schedule.DelayedCommand.after(0, target)
        sched.add(cmd)
        sched.run_pending()
        target.assert_called_once()
        assert not sched.queue

    def test_callback_scheduler(self):
        callback = mock.MagicMock()
        sched = schedule.CallbackScheduler(callback)
        target = mock.MagicMock()
        cmd = schedule.DelayedCommand.after(0, target)
        sched.add(cmd)
        sched.run_pending()
        callback.assert_called_once_with(target)

    def test_periodic_command(self):
        sched = schedule.InvokeScheduler()
        target = mock.MagicMock()
        cmd = schedule.PeriodicCommand.at_time(time.time() + 0.1, target)
        sched.add(cmd)
        sched.run_pending()
        target.assert_not_called()
        time.sleep(0.1)
        sched.run_pending()
        assert sched.queue
        target.assert_called_once()
        time.sleep(0.1)
        sched.run_pending()
        assert target.call_count == 2
<|endoftext|>


<commit_before>"""
Command function to merge multiple HXL datasets.
David Megginson
November 2014

Can use a whitelist of HXL tags, a blacklist, or both.

License: Public Domain
Documentation: http://hxlstandard.org
"""

import sys
import argparse
from copy import copy
from hxl.model import HXLSource, HXLColumn
from hxl.parser import HXLReader, writeHXL
from hxl.filters import parse_tags

class HXLMergeFilter(HXLSource):

    def __init__(self, source, merge_source, keys, tags):
        self.source = source
        self.merge_source = merge_source
        self.keys = keys
        self.merge_tags = tags
        self.saved_columns = None
        self.merge_map = None
        self.empty_result = [''] * len(tags)

    @property
    def columns(self):
        if self.saved_columns is None:
            self.saved_columns = self.source.columns + map(lambda tag: HXLColumn(hxlTag=tag), self.merge_tags)
        return self.saved_columns

    def next(self):
        if self.merge_map is None:
            self.merge_map = self._read_merge()
        row = copy(self.source.next())
        merge_values = self.merge_map.get(self._make_key(row))
        if not merge_values:
            merge_values = self.empty_result
        row.values = row.values + merge_values
        return row

    def _make_key(self, row):
        """
        Make a tuple key for a row.
        """
        values = []
        for key in self.keys:
            values.append(row.get(key))
        return tuple(values)

    def _read_merge(self):
        merge_map = {}
        for row in self.merge_source:
            values = []
            for tag in self.merge_tags:
                values.append(row.get(tag))
            merge_map[self._make_key(row)] = values
        return merge_map

#
# Command-line support
#

def run(args, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr):
    """
    Run hxlmerge with command-line arguments.
    @param args A list of arguments, excluding the script name
    @param stdin Standard input for the script
    @param stdout Standard output for the script
    @param stderr Standard error for the script
    """

    parser = argparse.ArgumentParser(description = 'Merge part of one HXL dataset into another.')
    parser.add_argument(
        'infile',
        help='HXL file to read (if omitted, use standard input).',
        nargs='?',
        type=argparse.FileType('r'),
        default=stdin
        )
    parser.add_argument(
        'outfile',
        help='HXL file to write (if omitted, use standard output).',
        nargs='?',
        type=argparse.FileType('w'),
        default=stdout
        )
    parser.add_argument(
        '-m',
        '--merge',
        help='HXL file to write (if omitted, use standard output).',
        metavar='filename',
        required=True,
        type=argparse.FileType('r')
        )
    parser.add_argument(
        '-k',
        '--keys',
        help='HXL tag(s) to use as a shared key.',
        metavar='tag,tag...',
        required=True,
        type=parse_tags
        )
    parser.add_argument(
        '-t',
        '--tags',
        help='Comma-separated list of column tags to include from the merge dataset.',
        metavar='tag,tag...',
        required=True,
        type=parse_tags
        )
    args = parser.parse_args(args)

    source = HXLReader(args.infile)
    filter = HXLMergeFilter(source, HXLReader(args.merge), args.keys, args.tags)
    writeHXL(args.outfile, filter)

# end
<commit_msg>Added better documentation comments.<commit_after>"""
Command function to merge multiple HXL datasets.
David Megginson
November 2014

Can use a whitelist of HXL tags, a blacklist, or both.

License: Public Domain
Documentation: http://hxlstandard.org
"""

import sys
import argparse
from copy import copy
from hxl.model import HXLSource, HXLColumn
from hxl.parser import HXLReader, writeHXL
from hxl.filters import parse_tags

class HXLMergeFilter(HXLSource):
    """
    Composable filter class to merge values from two HXL datasets.

    This is the class supporting the hxlmerge command-line utility.

    Warning: this filter may store a large amount of data in memory, depending on the merge.

    Because this class is a {@link hxl.model.HXLSource}, you can use
    it as the source to an instance of another filter class to build a
    dynamic, single-threaded processing pipeline.

    Usage:

    <pre>
    source = HXLReader(sys.stdin)
    merge_source = HXLReader(open('file-to-merge.csv', 'r'))
    filter = HXLMergeFilter(source, merge_source=merge_source, keys=['adm1_id'], tags=['adm1'])
    writeHXL(sys.stdout, filter)
    </pre>
    """

    def __init__(self, source, merge_source, keys, tags):
        """
        Constructor.
        @param source the HXL data source.
        @param merge_source a second HXL data source to merge into the first.
        @param keys the shared key hashtags to use for the merge
        @param tags the tags to include from the second dataset
        """
        self.source = source
        self.merge_source = merge_source
        self.keys = keys
        self.merge_tags = tags
        self.saved_columns = None
        self.merge_map = None
        self.empty_result = [''] * len(tags)

    @property
    def columns(self):
        """
        @return column definitions for the merged dataset
        """
        if self.saved_columns is None:
            self.saved_columns = self.source.columns + map(lambda tag: HXLColumn(hxlTag=tag), self.merge_tags)
        return self.saved_columns

    def next(self):
        """
        @return the next merged row of data
        """
        if self.merge_map is None:
            self.merge_map = self._read_merge()
        row = copy(self.source.next())
        merge_values = self.merge_map.get(self._make_key(row))
        if not merge_values:
            merge_values = self.empty_result
        row.values = row.values + merge_values
        return row

    def _make_key(self, row):
        """
        Make a tuple key for a row.
        """
        values = []
        for key in self.keys:
            values.append(row.get(key))
        return tuple(values)

    def _read_merge(self):
        """
        Read the second (merging) dataset into memory.
        Stores only the values necessary for the merge.
        @return a map of merge values
        """
        merge_map = {}
        for row in self.merge_source:
            values = []
            for tag in self.merge_tags:
                values.append(row.get(tag))
            merge_map[self._make_key(row)] = values
        return merge_map

#
# Command-line support
#

def run(args, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr):
    """
    Run hxlmerge with command-line arguments.
    @param args A list of arguments, excluding the script name
    @param stdin Standard input for the script
    @param stdout Standard output for the script
    @param stderr Standard error for the script
    """

    parser = argparse.ArgumentParser(description = 'Merge part of one HXL dataset into another.')
    parser.add_argument(
        'infile',
        help='HXL file to read (if omitted, use standard input).',
        nargs='?',
        type=argparse.FileType('r'),
        default=stdin
        )
    parser.add_argument(
        'outfile',
        help='HXL file to write (if omitted, use standard output).',
        nargs='?',
        type=argparse.FileType('w'),
        default=stdout
        )
    parser.add_argument(
        '-m',
        '--merge',
        help='HXL file to write (if omitted, use standard output).',
        metavar='filename',
        required=True,
        type=argparse.FileType('r')
        )
    parser.add_argument(
        '-k',
        '--keys',
        help='HXL tag(s) to use as a shared key.',
        metavar='tag,tag...',
        required=True,
        type=parse_tags
        )
    parser.add_argument(
        '-t',
        '--tags',
        help='Comma-separated list of column tags to include from the merge dataset.',
        metavar='tag,tag...',
        required=True,
        type=parse_tags
        )
    args = parser.parse_args(args)

    source = HXLReader(args.infile)
    filter = HXLMergeFilter(source, HXLReader(args.merge), args.keys, args.tags)
    writeHXL(args.outfile, filter)

# end
<|endoftext|>


<commit_before>import logging
import itertools
import h5py
import os
import numpy as np
import networkx as nx
from hytra.pluginsystem.plugin_manager import TrackingPluginManager
import hytra.core.probabilitygenerator as probabilitygenerator
import hytra.core.jsongraph
from hytra.core.jsongraph import negLog, listify, JsonTrackingGraph

"Anpassen aller DOOOOCS"

def getLogger():
    ''' logger to be used in this module '''
    return logging.getLogger(__name__)

class GapCloser(object):
    """
    Base class for all gap closing implementations. Use one of the derived classes
    that handle reading/writing data to the respective sources.
    """

    def __init__(self, pluginPaths=[os.path.abspath('../hytra/plugins')], verbose=False):
        self.Graph = None
        self.pluginManager = TrackingPluginManager(verbose=verbose, pluginPaths=pluginPaths)

        # should be filled by constructors of derived classes!
        self.model = None
        self.result = None

    def _determineTerminatingStartingTracks(self, input_ctc):
        """
        ** returns ** a list of detections (frame, id) sorted by frames to take into consideration for
        gap closing, which are tracks that end before time and ones that begin in the middle of the video.
        """
        # ctc format is: (id, start_frame, end_frame, parent_id)
        with open(input_ctc) as txt:
            content = txt.readlines()
        track_info = []
        for line in content:
            strings = line.strip('\n').split()
            track_info.append(tuple([int(t) for t in strings]))

        lookAt = []
        lastframe = max(track_info, key=lambda item:item[2])[2]
        parents = [t[-1] for t in track_info]

        for track in track_info:
            idx, start_frame, end_frame, parent_id = track
            if start_frame !=0 and parent_id == 0: # track started, but not at the beginning of the video, dah!
                lookAt.append((start_frame, idx, 'starting'))
            if end_frame != lastframe and idx not in parents:
                lookAt.append((end_frame, idx, 'ending')) # track ends AND has no continuation!

        return sorted(lookAt, key=lambda item:item[0])

    def _createGraph(self, lookAt):
        """
        Set up a networkx graph consisting of
            - last detection in tracks that end before the video does
            - first detection of tracks beginning in the middle of the video

        ** returns ** the `unGraph`
        """
        self.Graph = nx.DiGraph()
        def source(timestep, link):
            return int(timestep) - 1, link[0]

        def target(timestep, link):
            return int(timestep), link[1]

        def addNode(node):
            ''' add a node to the unresolved graph'''
            self.Graph.add_node((node[0], node[1]))

        # add nodes
        for node in lookAt:
            # find "frame-neighbours" to add connections to them
            for neighb in lookAt:
                # print node, neighb
                if node != neighb and node[0] == neighb[0] - 2 and node[2] == 'ending' and neighb[2] == 'starting':
                    addNode(node)
                    addNode(neighb)
                    self.Graph.add_edge((node[0], node[1]), (neighb[0], neighb[1]))
                    print node, neighb, "added edges"

        return self.Graph

    def _readLabelImage(self, timeframe):
        '''
        Should return the labelimage for the given timeframe
        '''
        raise NotImplementedError()

    def _minCostMaxFlowGapClosing(self, objectFeatures, transitionClassifier=None, transitionParameter=5.0):
        """
        Find the optimal assignments within the `Graph` by running min-cost max-flow from the
        `dpct` module.

        Converts the `Graph` to our JSON model structure, predicts the transition probabilities
        either using the given transitionClassifier, or using distance-based probabilities.

        **returns** a `nodeFlowMap` and `arcFlowMap` holding information on the usage of the respective nodes and links

        **Note:** cannot use `networkx` flow methods because they don't work with floating point weights.
        """

        map_uuid_to_ctc = {}
        trackingGraph = JsonTrackingGraph()
        for node in self.Graph.nodes_iter():
            additionalFeatures = {}
            if len(self.Graph.in_edges(node)) == 0:
                additionalFeatures['appearanceFeatures'] = [[0], [0]]
            if len(self.Graph.out_edges(node)) == 0:
                additionalFeatures['disappearanceFeatures'] = [[0], [0]]
            uuid = trackingGraph.addDetectionHypotheses([[0], [1]], **additionalFeatures)
            self.Graph.node[node]['id'] = uuid
            map_uuid_to_ctc[uuid] = node

        for edge in self.Graph.edges_iter():
            src = self.Graph.node[edge[0]]['id']
            dest = self.Graph.node[edge[1]]['id']

            featuresAtSrc = objectFeatures[edge[0]]
            featuresAtDest = objectFeatures[edge[1]]

            if transitionClassifier is not None:
                try:
                    featVec = self.pluginManager.applyTransitionFeatureVectorConstructionPlugins(
                        featuresAtSrc, featuresAtDest, transitionClassifier.selectedFeatures)
                except:
                    getLogger().error("Could not compute transition features of link {}->{}:".format(src, dest))
                    getLogger().error(featuresAtSrc)
                    getLogger().error(featuresAtDest)
                    raise
                featVec = np.expand_dims(np.array(featVec), axis=0)
                probs = transitionClassifier.predictProbabilities(featVec)[0]
            else:
                dist = np.linalg.norm(featuresAtDest['RegionCenter'] - featuresAtSrc['RegionCenter'])
                prob = np.exp(-dist / transitionParameter)
                probs = [1.0 - prob, prob]

            trackingGraph.addLinkingHypotheses(src, dest, listify(negLog(probs)))

        # track
        import dpct
        weights = {"weights": [1, 1, 1, 1]}
        mergerResult = dpct.trackMaxFlow(trackingGraph.model, weights)

        # transform results to dictionaries that can be indexed by id or (src,dest)
        nodeFlowMap = dict([(int(d['id']), int(d['value'])) for d in mergerResult['detectionResults']])
        arcFlowMap = dict([((int(l['src']), int(l['dest'])), int(l['value'])) for l in mergerResult['linkingResults']])

        ctc_arcFlowMap = {}
        # translate edges back to ctc format
        for edge in arcFlowMap.keys():
            src = map_uuid_to_ctc[edge[0]]
            dest = map_uuid_to_ctc[edge[1]]
            ctc_arcFlowMap[(src, dest)] = 1 # poate altfel aici, lista fara value, care oricum e unu. Sau?

        # return nodeFlowMap, arcFlowMap
        return ctc_arcFlowMap

    def _saveTracks(self, ctc_arcFlowMap, input_ctc, output_ctc):
        """
        Update the ctc result format (res_track.txt) by adding the right parents to the closed gaps.
        So read in the old sresult txtx file and write a new one.
        """

        with open(input_ctc) as input_txt:
            content = input_txt.readlines()

        with open(output_ctc, 'wt') as f:
            for line in content:
                for edge in ctc_arcFlowMap:
                    strings = line.strip('\n').split()
                    track_info = tuple([int(t) for t in strings])
                    idx, start_frame, end_frame, parent_id = track_info
                    if edge[1][1] == idx:
                        # if idx == 13: # 10
                        #     continue
                        print line
                        parent_id = edge[0][1]
                        line ="{} {} {} {}\n".format(idx, start_frame, end_frame, parent_id)
                        print line
                f.write(line)

    def _computeObjectFeatures(self, labelImages):
        '''
        Return the features per object as nested dictionaries:
        { (int(Timestep), int(Id)):{ "FeatureName" : np.array(value), "NextFeature": ...} }
        '''
        pass

    # ------------------------------------------------------------
    def run(self, input_ctc, output_ctc, transition_classifier_filename=None, transition_classifier_path=None):
        """
        Run merger resolving

        1. find mergers in the given model and result
        2. build graph of the unresolved (merger) nodes and their direct neighbors
        3. use a mergerResolving plugin to refine the merger nodes and their segmentation
        4. run min-cost max-flow tracking to find the fate of all the de-merged objects
        5. export refined segmentation, update member variables `model` and `result`

        """

        traxelIdPerTimestepToUniqueIdMap, uuidToTraxelMap = hytra.core.jsongraph.getMappingsBetweenUUIDsAndTraxels(self.model)
        timesteps = [t for t in traxelIdPerTimestepToUniqueIdMap.keys()]

        mergers, detections, links, divisions = hytra.core.jsongraph.getMergersDetectionsLinksDivisions(self.result, uuidToTraxelMap)
        lookAt = self._determineTerminatingStartingTracks(input_ctc)
        print "LOOK AT", lookAt

        # ------------------------------------------------------------
        if lookAt == None:
            getLogger().info("There are no gaps to close, nothing to be done. Writing the output...")
            # segmentation
            labelImages = {}
            for t in timesteps:
                # use image provider plugin to load labelimage
                labelImage = self._readLabelImage(int(t))
                labelImages[t] = labelImage
        else:
            # set up unresolved graph and then refine the nodes to get the resolved graph
            self._createGraph(lookAt)
            labelImages = {}
            for t in timesteps:
                # use image provider plugin to load labelimage
                labelImage = self._readLabelImage(int(t))
                labelImages[t] = labelImage

            # ------------------------------------------------------------
            # compute new object features
            objectFeatures = self._computeObjectFeatures(labelImages)

            # ------------------------------------------------------------
            # load transition classifier if any
            if transition_classifier_filename is not None:
                getLogger().info("\tLoading transition classifier")
                transitionClassifier = probabilitygenerator.RandomForestClassifier(
                    transition_classifier_path, transition_classifier_filename)
            else:
                getLogger().info("\tUsing distance based transition energies")
                transitionClassifier = None

            # ------------------------------------------------------------
            # run min-cost max-flow to find gaps worthy to be closed
            getLogger().info("Running min-cost max-flow to find closed gaps assignments")

            ctcArcFlowMap = self._minCostMaxFlowGapClosing(objectFeatures, transitionClassifier)

            # ------------------------------------------------------------
            # fuse results into a new solution

            self.result = self._saveTracks(ctcArcFlowMap, input_ctc, output_ctc)
<commit_msg>Only probability-based gap closing. Tested only on Fluo-SIM 02<commit_after>import logging
import itertools
import h5py
import os
import numpy as np
import networkx as nx
from hytra.pluginsystem.plugin_manager import TrackingPluginManager
import hytra.core.probabilitygenerator as probabilitygenerator
import hytra.core.jsongraph
from hytra.core.jsongraph import negLog, listify, JsonTrackingGraph

"Anpassen aller DOOOOCS"

def getLogger():
    ''' logger to be used in this module '''
    return logging.getLogger(__name__)

class GapCloser(object):
    """
    Base class for all gap closing implementations. Use one of the derived classes
    that handle reading/writing data to the respective sources.
    """

    def __init__(self, pluginPaths=[os.path.abspath('../hytra/plugins')], verbose=False):
        self.Graph = None
        self.pluginManager = TrackingPluginManager(verbose=verbose, pluginPaths=pluginPaths)

        # should be filled by constructors of derived classes!
        self.model = None
        self.result = None

    def _determineTerminatingStartingTracks(self, input_ctc):
        """
        ** returns ** a list of detections (frame, id) sorted by frames to take into consideration for
        gap closing, which are tracks that end before time and ones that begin in the middle of the video.
        """
        # ctc format is: (id, start_frame, end_frame, parent_id)
        with open(input_ctc) as txt:
            content = txt.readlines()
        track_info = []
        for line in content:
            strings = line.strip('\n').split()
            track_info.append(tuple([int(t) for t in strings]))

        lookAt = []
        lastframe = max(track_info, key=lambda item:item[2])[2]
        parents = [t[-1] for t in track_info]

        for track in track_info:
            idx, start_frame, end_frame, parent_id = track
            if start_frame !=0 and parent_id == 0: # track started, but not at the beginning of the video, dah!
                lookAt.append((start_frame, idx, 'starting'))
            if end_frame != lastframe and idx not in parents:
                lookAt.append((end_frame, idx, 'ending')) # track ends AND has no continuation!

        return sorted(lookAt, key=lambda item:item[0])

    def _createGraph(self, lookAt):
        """
        Set up a networkx graph consisting of
            - last detection in tracks that end before the video does
            - first detection of tracks beginning in the middle of the video

        ** returns ** the `unGraph`
        """
        self.Graph = nx.DiGraph()
        def source(timestep, link):
            return int(timestep) - 1, link[0]

        def target(timestep, link):
            return int(timestep), link[1]

        def addNode(node):
            ''' add a node to the unresolved graph'''
            self.Graph.add_node((node[0], node[1]))

        # add nodes
        for node in lookAt:
            # find "frame-neighbours" to add connections to them
            for neighb in lookAt:
                # print node, neighb
                if node != neighb and node[0] == neighb[0] - 2 and node[2] == 'ending' and neighb[2] == 'starting':
                    addNode(node)
                    addNode(neighb)
                    self.Graph.add_edge((node[0], node[1]), (neighb[0], neighb[1]))
                    print node, neighb, "added edges"

        return self.Graph

    def _readLabelImage(self, timeframe):
        '''
        Should return the labelimage for the given timeframe
        '''
        raise NotImplementedError()

    def _minCostMaxFlowGapClosing(self, objectFeatures, transitionClassifier=None, transitionParameter=5.0, treshold=0.15):
        """
        Find the optimal assignments within the `Graph` by running min-cost max-flow from the
        `dpct` module.

        Converts the `Graph` to our JSON model structure, predicts the transition probabilities
        either using the given transitionClassifier, or using distance-based probabilities.

        **returns** a `nodeFlowMap` and `arcFlowMap` holding information on the usage of the respective nodes and links

        **Note:** cannot use `networkx` flow methods because they don't work with floating point weights.
        """

        map_uuid_to_ctc = {}
        trackingGraph = JsonTrackingGraph()
        for node in self.Graph.nodes_iter():
            additionalFeatures = {}
            if len(self.Graph.in_edges(node)) == 0:
                additionalFeatures['appearanceFeatures'] = [[0], [0]]
            if len(self.Graph.out_edges(node)) == 0:
                additionalFeatures['disappearanceFeatures'] = [[0], [0]]
            uuid = trackingGraph.addDetectionHypotheses([[0], [1]], **additionalFeatures)
            self.Graph.node[node]['id'] = uuid
            map_uuid_to_ctc[uuid] = node

        ctc_arcFlowMap = {}
        for edge in self.Graph.edges_iter():
            src = self.Graph.node[edge[0]]['id']
            dest = self.Graph.node[edge[1]]['id']

            featuresAtSrc = objectFeatures[edge[0]]
            featuresAtDest = objectFeatures[edge[1]]

            if transitionClassifier is not None:
                try:
                    featVec = self.pluginManager.applyTransitionFeatureVectorConstructionPlugins(
                        featuresAtSrc, featuresAtDest, transitionClassifier.selectedFeatures)
                except:
                    getLogger().error("Could not compute transition features of link {}->{}:".format(src, dest))
                    getLogger().error(featuresAtSrc)
                    getLogger().error(featuresAtDest)
                    raise
                featVec = np.expand_dims(np.array(featVec), axis=0)
                probs = transitionClassifier.predictProbabilities(featVec)[0]
            else:
                dist = np.linalg.norm(featuresAtDest['RegionCenter'] - featuresAtSrc['RegionCenter'])
                prob = np.exp(-dist / transitionParameter)
                probs = [1.0 - prob, prob]

            print probs, edge

            if probs[1] > treshold:
                ctc_arcFlowMap[edge] = 1
            else:
                ctc_arcFlowMap[edge] = 0

            # trackingGraph.addLinkingHypotheses(src, dest, listify(negLog(probs)))

        # # track
        # import dpct
        # weights = {"weights": [1, 1, 1, 1]}
        # # mergerResult = dpct.trackMaxFlow(trackingGraph.model, weights)
        # mergerResult = dpct.trackFlowBased(trackingGraph.model, weights)

        # # transform results to dictionaries that can be indexed by id or (src,dest)
        # nodeFlowMap = dict([(int(d['id']), int(d['value'])) for d in mergerResult['detectionResults']])
        # arcFlowMap = dict([((int(l['src']), int(l['dest'])), int(l['value'])) for l in mergerResult['linkingResults']])

        # ctc_arcFlowMap = {}
        # # translate edges back to ctc format
        # for edge in arcFlowMap.keys():
        #     src = map_uuid_to_ctc[edge[0]]
        #     dest = map_uuid_to_ctc[edge[1]]
        #     ctc_arcFlowMap[(src, dest)] = 1 # poate altfel aici, lista fara value, care oricum e unu. Sau?

        # # return nodeFlowMap, arcFlowMap
        # print arcFlowMap
        return ctc_arcFlowMap

    def _saveTracks(self, ctc_arcFlowMap, input_ctc, output_ctc):
        """
        Update the ctc result format (res_track.txt) by adding the right parents to the closed gaps.
        So read in the old sresult txtx file and write a new one.
        """

        if ctc_arcFlowMap is not None:
            with open(input_ctc) as input_txt:
                content = input_txt.readlines()

            with open(output_ctc, 'wt') as f:
                for line in content:
                    for edge in ctc_arcFlowMap:
                        strings = line.strip('\n').split()
                        track_info = tuple([int(t) for t in strings])
                        idx, start_frame, end_frame, parent_id = track_info
                        if edge[1][1] == idx and ctc_arcFlowMap[edge] == 1:
                            parent_id = edge[0][1]
                            line ="{} {} {} {}\n".format(idx, start_frame, end_frame, parent_id)
                    f.write(line)

    def _computeObjectFeatures(self, labelImages):
        '''
        Return the features per object as nested dictionaries:
        { (int(Timestep), int(Id)):{ "FeatureName" : np.array(value), "NextFeature": ...} }
        '''
        pass

    # ------------------------------------------------------------
    def run(self, input_ctc, output_ctc, transition_classifier_filename=None, transition_classifier_path=None):
        """
        Run merger resolving

        1. find mergers in the given model and result
        2. build graph of the unresolved (merger) nodes and their direct neighbors
        3. use a mergerResolving plugin to refine the merger nodes and their segmentation
        4. run min-cost max-flow tracking to find the fate of all the de-merged objects
        5. export refined segmentation, update member variables `model` and `result`

        """

        traxelIdPerTimestepToUniqueIdMap, uuidToTraxelMap = hytra.core.jsongraph.getMappingsBetweenUUIDsAndTraxels(self.model)
        timesteps = [t for t in traxelIdPerTimestepToUniqueIdMap.keys()]

        mergers, detections, links, divisions = hytra.core.jsongraph.getMergersDetectionsLinksDivisions(self.result, uuidToTraxelMap)
        lookAt = self._determineTerminatingStartingTracks(input_ctc)

        # ------------------------------------------------------------
        if lookAt == None:
            getLogger().info("There are no gaps to close, nothing to be done. Writing the output...")
            # segmentation
            labelImages = {}
            for t in timesteps:
                # use image provider plugin to load labelimage
                labelImage = self._readLabelImage(int(t))
                labelImages[t] = labelImage
        else:
            # set up unresolved graph and then refine the nodes to get the resolved graph
            self._createGraph(lookAt)
            labelImages = {}
            for t in timesteps:
                # use image provider plugin to load labelimage
                labelImage = self._readLabelImage(int(t))
                labelImages[t] = labelImage

            # ------------------------------------------------------------
            # compute new object features
            objectFeatures = self._computeObjectFeatures(labelImages)

            # ------------------------------------------------------------
            # load transition classifier if any
            if transition_classifier_filename is not None:
                getLogger().info("\tLoading transition classifier")
                transitionClassifier = probabilitygenerator.RandomForestClassifier(
                    transition_classifier_path, transition_classifier_filename)
            else:
                getLogger().info("\tUsing distance based transition energies")
                transitionClassifier = None

            # run min-cost max-flow to find gaps worthy to be closed
            getLogger().info("Running min-cost max-flow to find closed gaps assignments")

            ctcArcFlowMap = self._minCostMaxFlowGapClosing(objectFeatures, transitionClassifier)

            # fuse results into a new solution
            self.result = self._saveTracks(ctcArcFlowMap, input_ctc, output_ctc)
<|endoftext|>


<commit_before># Copyright 2019 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generic TFX PrestoExampleGen executor."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime

import apache_beam as beam
import prestodb
from proto import presto_config_pb2
import tensorflow as tf
from typing import Any, Dict, Iterable, List, Text, Tuple

from google.protobuf import json_format
from tfx import types
from tfx.components.example_gen import base_example_gen_executor
from tfx.proto import example_gen_pb2

@beam.typehints.with_input_types(Text)
@beam.typehints.with_output_types(beam.typehints.Iterable[Tuple[Text, Text,
                                                                Any]])
class _ReadPrestoDoFn(beam.DoFn):
  """Beam DoFn class that reads from Presto.

  Attributes:
    cursor: A prestodb.dbapi.Cursor object that reads records from Presto table.
  """

  def __init__(self, client: prestodb.dbapi.Connection):
    self.cursor = client.cursor()

  def process(self, query: Text) -> Iterable[Tuple[Text, Text, Any]]:
    """Yields rows from query results.

    Args:
      query: A SQL query used to return results from Presto table.

    Yields:
      One row from the query result, represented by a list of tuples. Each tuple
      contains information on column name, column data type, data.
    """
    self.cursor.execute(query)
    rows = self.cursor.fetchall()
    if rows:
      cols = []
      col_types = []
      # Returns a list of (column_name, column_type, None, ...)
      # https://github.com/prestodb/presto-python-client/blob/master/prestodb/dbapi.py#L199
      for metadata in self.cursor.description:
        cols.append(metadata[0])
        col_types.append(metadata[1])

      for r in rows:
        yield zip(cols, col_types, r)

  def teardown(self):
    if self.cursor:
      self.cursor.close()

def _deserialize_conn_config(
    conn_config: presto_config_pb2.PrestoConnConfig
) -> prestodb.dbapi.Connection:
  """Deserializes Presto connection config to Presto client.

  Args:
    conn_config: Protobuf-encoded connection config for Presto client.

  Returns:
    A prestodb.dbapi.Connection instance initialized with user-supplied
    parameters.
  """
  params = {'host': conn_config.host}  # Required field
  # Only deserialize rest of parameters if set by user
  if conn_config.HasField('port'):
    params['port'] = conn_config.port
  if conn_config.HasField('user'):
    params['user'] = conn_config.user
  if conn_config.HasField('source'):
    params['source'] = conn_config.source
  if conn_config.HasField('catalog'):
    params['catalog'] = conn_config.catalog
  if conn_config.HasField('schema'):
    params['schema'] = conn_config.schema
  if conn_config.HasField('http_scheme'):
    params['http_scheme'] = conn_config.http_scheme
  if conn_config.WhichOneof('opt_auth'):
    params['auth'] = _deserialize_auth_config(conn_config)
  if conn_config.HasField('max_attempts'):
    params['max_attempts'] = conn_config.max_attempts
  if conn_config.HasField('request_timeout'):
    params['request_timeout'] = conn_config.request_timeout

  return prestodb.dbapi.connect(**params)

def _deserialize_auth_config(
    conn_config: presto_config_pb2.PrestoConnConfig
) -> prestodb.auth.Authentication:
  """Extracts from conn config the deserialized Presto Authentication class.

  Args:
    conn_config: Protobuf-encoded connection config for Presto client.

  Returns:
    A prestodb.auth.Authentication instance initialized with user-supplied
    parameters.

  Raises:
    RuntimeError: if authentication type is not currently supported.
  """
  if conn_config.HasField('basic_auth'):
    return prestodb.auth.BasicAuthentication(conn_config.basic_auth.username,
                                             conn_config.basic_auth.password)
    # TODO(actam): Support KerberosAuth.
  else:
    raise RuntimeError('Authentication type not supported.')

def _row_to_example(
    instance: Iterable[Tuple[Text, Text, Any]]) -> tf.train.Example:
  """Convert presto result row to tf example."""
  feature = {}
  for key, data_type, value in instance:
    if value is None:
      feature[key] = tf.train.Feature()
    elif data_type in {'tinyint', 'smallint', 'integer', 'bigint'}:
      feature[key] = tf.train.Feature(
          int64_list=tf.train.Int64List(value=[value]))
    elif data_type in {'real', 'double', 'decimal'}:
      feature[key] = tf.train.Feature(
          float_list=tf.train.FloatList(value=[value]))
    elif data_type in {'varchar', 'char'}:
      feature[key] = tf.train.Feature(
          bytes_list=tf.train.BytesList(value=[tf.compat.as_bytes(value)]))
    elif data_type in {'timestamp'}:
      value = int(datetime.datetime.fromisoformat(value).timestamp())
      feature[key] = tf.train.Feature(
          int64_list=tf.train.Int64List(value=[value]))
    else:
      # TODO(actam): support more types
      # https://prestodb.github.io/docs/current/language/types
      raise RuntimeError(
          'Presto column type {} is not supported.'.format(data_type))
  return tf.train.Example(features=tf.train.Features(feature=feature))

@beam.ptransform_fn
@beam.typehints.with_input_types(beam.Pipeline)
@beam.typehints.with_output_types(tf.train.Example)
def _PrestoToExample(  # pylint: disable=invalid-name
    pipeline: beam.Pipeline,
    input_dict: Dict[Text, List[types.Artifact]],  # pylint: disable=unused-argument
    exec_properties: Dict[Text, Any],
    split_pattern: Text) -> beam.pvalue.PCollection:
  """Read from Presto and transform to TF examples.

  Args:
    pipeline: beam pipeline.
    input_dict: Input dict from input key to a list of Artifacts.
    exec_properties: A dict of execution properties.
    split_pattern: Split.pattern in Input config, a Presto sql string.

  Returns:
    PCollection of TF examples.
  """
  conn_config = example_gen_pb2.CustomConfig()
  json_format.Parse(exec_properties['custom_config'], conn_config)
  presto_config = presto_config_pb2.PrestoConnConfig()
  conn_config.custom_config.Unpack(presto_config)

  client = _deserialize_conn_config(presto_config)
  return (pipeline
          | 'Query' >> beam.Create([split_pattern])
          | 'QueryTable' >> beam.ParDo(_ReadPrestoDoFn(client))
          | 'ToTFExample' >> beam.Map(_row_to_example))

class Executor(base_example_gen_executor.BaseExampleGenExecutor):
  """Generic TFX PrestoExampleGen executor."""

  def GetInputSourceToExamplePTransform(self) -> beam.PTransform:
    """Returns PTransform for Presto to TF examples."""
    return _PrestoToExample
<commit_msg>Move todo to bug id instead.<commit_after># Copyright 2019 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generic TFX PrestoExampleGen executor."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime

import apache_beam as beam
import prestodb
from proto import presto_config_pb2
import tensorflow as tf
from typing import Any, Dict, Iterable, List, Text, Tuple

from google.protobuf import json_format
from tfx import types
from tfx.components.example_gen import base_example_gen_executor
from tfx.proto import example_gen_pb2

@beam.typehints.with_input_types(Text)
@beam.typehints.with_output_types(beam.typehints.Iterable[Tuple[Text, Text,
                                                                Any]])
class _ReadPrestoDoFn(beam.DoFn):
  """Beam DoFn class that reads from Presto.

  Attributes:
    cursor: A prestodb.dbapi.Cursor object that reads records from Presto table.
  """

  def __init__(self, client: prestodb.dbapi.Connection):
    self.cursor = client.cursor()

  def process(self, query: Text) -> Iterable[Tuple[Text, Text, Any]]:
    """Yields rows from query results.

    Args:
      query: A SQL query used to return results from Presto table.

    Yields:
      One row from the query result, represented by a list of tuples. Each tuple
      contains information on column name, column data type, data.
    """
    self.cursor.execute(query)
    rows = self.cursor.fetchall()
    if rows:
      cols = []
      col_types = []
      # Returns a list of (column_name, column_type, None, ...)
      # https://github.com/prestodb/presto-python-client/blob/master/prestodb/dbapi.py#L199
      for metadata in self.cursor.description:
        cols.append(metadata[0])
        col_types.append(metadata[1])

      for r in rows:
        yield zip(cols, col_types, r)

  def teardown(self):
    if self.cursor:
      self.cursor.close()

def _deserialize_conn_config(
    conn_config: presto_config_pb2.PrestoConnConfig
) -> prestodb.dbapi.Connection:
  """Deserializes Presto connection config to Presto client.

  Args:
    conn_config: Protobuf-encoded connection config for Presto client.

  Returns:
    A prestodb.dbapi.Connection instance initialized with user-supplied
    parameters.
  """
  params = {'host': conn_config.host}  # Required field
  # Only deserialize rest of parameters if set by user
  if conn_config.HasField('port'):
    params['port'] = conn_config.port
  if conn_config.HasField('user'):
    params['user'] = conn_config.user
  if conn_config.HasField('source'):
    params['source'] = conn_config.source
  if conn_config.HasField('catalog'):
    params['catalog'] = conn_config.catalog
  if conn_config.HasField('schema'):
    params['schema'] = conn_config.schema
  if conn_config.HasField('http_scheme'):
    params['http_scheme'] = conn_config.http_scheme
  if conn_config.WhichOneof('opt_auth'):
    params['auth'] = _deserialize_auth_config(conn_config)
  if conn_config.HasField('max_attempts'):
    params['max_attempts'] = conn_config.max_attempts
  if conn_config.HasField('request_timeout'):
    params['request_timeout'] = conn_config.request_timeout

  return prestodb.dbapi.connect(**params)

def _deserialize_auth_config(
    conn_config: presto_config_pb2.PrestoConnConfig
) -> prestodb.auth.Authentication:
  """Extracts from conn config the deserialized Presto Authentication class.

  Args:
    conn_config: Protobuf-encoded connection config for Presto client.

  Returns:
    A prestodb.auth.Authentication instance initialized with user-supplied
    parameters.

  Raises:
    RuntimeError: if authentication type is not currently supported.
  """
  if conn_config.HasField('basic_auth'):
    return prestodb.auth.BasicAuthentication(conn_config.basic_auth.username,
                                             conn_config.basic_auth.password)
    # TODO(b/140266796): Support KerberosAuth.
  else:
    raise RuntimeError('Authentication type not supported.')

def _row_to_example(
    instance: Iterable[Tuple[Text, Text, Any]]) -> tf.train.Example:
  """Convert presto result row to tf example."""
  feature = {}
  for key, data_type, value in instance:
    if value is None:
      feature[key] = tf.train.Feature()
    elif data_type in {'tinyint', 'smallint', 'integer', 'bigint'}:
      feature[key] = tf.train.Feature(
          int64_list=tf.train.Int64List(value=[value]))
    elif data_type in {'real', 'double', 'decimal'}:
      feature[key] = tf.train.Feature(
          float_list=tf.train.FloatList(value=[value]))
    elif data_type in {'varchar', 'char'}:
      feature[key] = tf.train.Feature(
          bytes_list=tf.train.BytesList(value=[tf.compat.as_bytes(value)]))
    elif data_type in {'timestamp'}:
      value = int(datetime.datetime.fromisoformat(value).timestamp())
      feature[key] = tf.train.Feature(
          int64_list=tf.train.Int64List(value=[value]))
    else:
      # TODO(b/140266796): support more types
      # https://prestodb.github.io/docs/current/language/types
      raise RuntimeError(
          'Presto column type {} is not supported.'.format(data_type))
  return tf.train.Example(features=tf.train.Features(feature=feature))

@beam.ptransform_fn
@beam.typehints.with_input_types(beam.Pipeline)
@beam.typehints.with_output_types(tf.train.Example)
def _PrestoToExample(  # pylint: disable=invalid-name
    pipeline: beam.Pipeline,
    input_dict: Dict[Text, List[types.Artifact]],  # pylint: disable=unused-argument
    exec_properties: Dict[Text, Any],
    split_pattern: Text) -> beam.pvalue.PCollection:
  """Read from Presto and transform to TF examples.

  Args:
    pipeline: beam pipeline.
    input_dict: Input dict from input key to a list of Artifacts.
    exec_properties: A dict of execution properties.
    split_pattern: Split.pattern in Input config, a Presto sql string.

  Returns:
    PCollection of TF examples.
  """
  conn_config = example_gen_pb2.CustomConfig()
  json_format.Parse(exec_properties['custom_config'], conn_config)
  presto_config = presto_config_pb2.PrestoConnConfig()
  conn_config.custom_config.Unpack(presto_config)

  client = _deserialize_conn_config(presto_config)
  return (pipeline
          | 'Query' >> beam.Create([split_pattern])
          | 'QueryTable' >> beam.ParDo(_ReadPrestoDoFn(client))
          | 'ToTFExample' >> beam.Map(_row_to_example))

class Executor(base_example_gen_executor.BaseExampleGenExecutor):
  """Generic TFX PrestoExampleGen executor."""

  def GetInputSourceToExamplePTransform(self) -> beam.PTransform:
    """Returns PTransform for Presto to TF examples."""
    return _PrestoToExample
<|endoftext|>