<commit_before>#!/usr/bin/env python3
import cv2
import numpy as np
import sys
from scipy.ndimage import rotate

def get_rotated_pos(x, y, rot):
    return (int(x * np.cos(rot) - y * np.sin(rot) + 0.5),
            int(x * np.sin(rot) + y * np.cos(rot) + 0.5))

bottom_right = lambda N, M, k: np.fliplr(np.tri(N, M, k=k)) == 1
top_left = lambda N, M, k: np.flipud(np.tri(N, M, k=k)) == 1
bottom_left = lambda N, M, k: np.tri(N, M, k=k) == 1
top_right = lambda N, M, k: np.flipud(np.fliplr(np.tri(N, M, k=k))) == 1

def find_connected_nodes(image, i, nodes):
    out = []
    x, y = nodes[i]
    for j in range(i + 1, len(nodes)):
        if i == j: continue
        distance = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[j]), 2)))
        dst_x, dst_y = nodes[j]

        angle = np.arctan2(dst_y - y, dst_x - x)

        min_x = min(x, dst_x)
        min_y = min(y, dst_y)
        max_x = max(x, dst_x)
        max_y = max(y, dst_y)
        border = 2
        cutout = np.copy(image[min_y - border : max_y + border + 1, min_x - border : max_x + border + 1])

        if ((x - min_x) == 0 and (y - min_y) == 0) or ((dst_x - min_x) == 0 and (dst_y - min_y) == 0):
            cutout[top_right(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
            cutout[bottom_left(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
        else:
            cutout[top_left(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
            cutout[bottom_right(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255

        contours = find_tree_contours(cutout, 2)
        for contour in contours:
            other = False
            for k in range(len(nodes)):
                if k == j or k == i: continue
                k_x, k_y = nodes[k]
                k_x = k_x - min_x + border
                k_y = k_y - min_y + border
                if k_x >= border and k_x < cutout.shape[-1] - border and k_y >= border and k_y < cutout.shape[-2] - border:
                    if contour[k_y, k_x] < 200:
                        other = True
                        break

            # Check if there's not something along the way
            if i == 43 and j == 57:
                print(i, j, other,
                      contour[y     - min_y + border,     x - min_x + border],
                      contour[dst_y - min_y + border, dst_x - min_x + border],
                      nodes[i], nodes[j])
                print(image[dst_y - 2 : dst_y + 3, dst_x - 2 : dst_x + 3])
                print(cutout[:5, :5])
                print(contour[:5, :5])

                cv2.imshow('1', cutout)
                cv2.imshow('2', contour)
                cv2.waitKey(0)

            if other == False and contour[y - min_y + border, x - min_x + border] < 200 and contour[dst_y - min_y + border, dst_x - min_x + border] < 200:
                out.append(j)
    return out

def find_nearest_connected_nodes(image, i, nodes):
    out = []
    x, y = nodes[i]
    min_j = i
    min_dist = np.max(image.shape)
    distances = []
    good_js = []
    for j in range(len(nodes)):
        distance = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[j]), 2)))
        distances.append(distance)
        if i == j: continue
        dst_x, dst_y = nodes[j]

        angle = np.arctan2(dst_y - y, dst_x - x)

        min_x = min(x, dst_x)
        min_y = min(y, dst_y)
        max_x = max(x, dst_x)
        max_y = max(y, dst_y)
        cutout = np.copy(image[min_y - 1  : max_y + 2, min_x - 1 : max_x + 2])

        contours = find_tree_contours(cutout, 1)
        for contour in contours:
            if contour[y - min_y + 1, x - min_x + 1] < 128 and contour[dst_y - min_y + 1, dst_x - min_x + 1] < 128:
                if distance < min_dist:
                    min_j = j
                    min_dist = distance
                good_js.append(j)
    for j in good_js:
        if distances[j] < min_dist * 1.5:
            out.append(nodes[j])
    return out

# Nodes from corners
def nodes_from_corners(image, gray, points, max_dist=16, iterations=1):
    if iterations == 0:
        out = []

        for j, point in enumerate(points):
            x, y = map(int, point)
            A = gray[y - 1 : y + 2, x - 1 : x + 2]
            off_y, off_x = np.array(np.unravel_index(A.argmin(), A.shape)) - 1
            new_x, new_y = x + off_x, y + off_y
            if gray[new_y, new_x] < gray[y, x]:
                out.append((new_x,new_y))
            else:
                out.append((x,y))

        return out
    print('iternations:', iterations)
    nodes = []
    while len(points):
        point = points[-1]
        points.pop()
        neighbors = [point]
        for i in reversed(range(len(points))):
            dist = np.sqrt(np.sum(np.power(np.array(point) - np.array(points[i]), 2)))
            if dist < max_dist:
                neighbors.append(points[i])
                points.pop(i)
        neighbors = np.array(neighbors)
        point = neighbors[np.random.randint(len(neighbors))]
        point = np.mean(neighbors, axis=0)
        print('Remaining corner points to assign a corner:', len(points))
        if iterations == 1:
            #image[point[1], point[0], :] = (255, 0, 0)
            cv2.circle(image, tuple(map(int, point)), 4, (255,128,0), 2)
        nodes.append(point)
    return nodes_from_corners(image, gray, nodes, max_dist + 2, iterations - 1)

# Find the contour of the tree
def find_tree_contour(gray, param, min_frac=0.5):
    _,thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV) # threshold
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    dilated = cv2.dilate(thresh,kernel,iterations = param) # dilate
    im2, contours, hierarchy = cv2.findContours(dilated, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    out = np.zeros(gray.shape, dtype=np.uint8) + 255

    for i, contour in enumerate(contours):
        [x,y,w,h] = cv2.boundingRect(contour)
        if gray.shape[-1] > gray.shape[-2]:
            if w < gray.shape[-1] * min_frac:
                continue
        elif gray.shape[-1] < gray.shape[-2]:
            if h < gray.shape[-2] * min_frac:
                continue
        else:
            if h < gray.shape[-2] * min_frac or w < gray.shape[-1] * min_frac:
                continue
        #cv2.rectangle(img, (x,y), (x+w,y+h), (255,0,255), 1)
        cv2.drawContours(out, contours, i, 0, -1)

    out2 = np.zeros(gray.shape, dtype=np.uint8) + 255
    out2[out == 0] = gray[out == 0]
    return out2

def find_tree_contours(gray, param):
    _,thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV) # threshold
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    dilated = cv2.dilate(thresh,kernel,iterations = param) # dilate
    im2, contours, hierarchy = cv2.findContours(dilated, cv2.RETR_LIST, cv2.CHAIN_APPROX_TC89_KCOS)

    outs = []
    for i, contour in enumerate(contours):
        [x,y,w,h] = cv2.boundingRect(contour)
        if h < gray.shape[0] * 0.5 or w < gray.shape[1] * 0.5: continue
        #cv2.rectangle(img, (x,y), (x+w,y+h), (255,0,255), 1)
        out = np.zeros(gray.shape, dtype=np.uint8) + 255
        cv2.drawContours(out, contours, i, 0, -1)
        outs.append(out)
    #out2[out < 2010] = 0
    return outs

def main():
    filename = sys.argv[1]
    param = int(sys.argv[2])

    img = cv2.imread(filename)
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
            img[y,x,:] = np.min(img[y,x])
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    tree_contour = find_tree_contour(gray, param)
    cv2.imshow('dst2', tree_contour)

    dst = cv2.cornerHarris(tree_contour, 3, 3, 0.01)
    dst = cv2.dilate(dst, None)
    corners = dst > 0.05 * dst.max()

    points = []
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
            if corners[y,x]:
                points.append((x,y))

    nodes = nodes_from_corners(img, gray, points, max_dist=8, iterations=3)
    for i, node in enumerate(nodes):
        x, y = node
        print(i, gray[y, x])

    # Figure out node neighbors
    node_neighbors = {}
    for i, node in enumerate(nodes):
        node_neighbors[i] = find_connected_nodes(gray, i, nodes)

    # Do reciprocal adding of nodes
    for i in range(len(nodes)):
        for j in node_neighbors[i]:
            if i not in node_neighbors[j]:
                node_neighbors[j].append(i)

    # Delete nodes with only 2 neighbors as they are connections:
    for i in range(len(nodes)):
        if len(node_neighbors[i]) == 2:
            n1, n2 = node_neighbors[i]
            for j in node_neighbors:
                print(node_neighbors[j], i, j)
                node_neighbors[j] = [x for x in node_neighbors[j] if x != i]
            node_neighbors[n1].append(n2)
            node_neighbors[n2].append(n1)
            del node_neighbors[i]

    # Print lines
    for i in sorted(node_neighbors):
        node_x, node_y = nodes[i]
        for neighbor in node_neighbors[i]:
            neighbor_x, neighbor_y = nodes[neighbor]
            cv2.line(img, (node_x, node_y), (neighbor_x, neighbor_y), (255,0,0), 2)

    for i in sorted(node_neighbors):
        print(i, node_neighbors[i])

    for i in sorted(node_neighbors):
        cv2.putText(img, str(i), nodes[i], cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))

    cv2.imshow('dst', img)
    if cv2.waitKey(0) & 0xff == 27:
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
<commit_msg>Added triangle remover... looks pretty good.<commit_after>#!/usr/bin/env python3
import cv2
import numpy as np
import sys
from scipy.ndimage import rotate

def get_rotated_pos(x, y, rot):
    return (int(x * np.cos(rot) - y * np.sin(rot) + 0.5),
            int(x * np.sin(rot) + y * np.cos(rot) + 0.5))

bottom_right = lambda N, M, k: np.fliplr(np.tri(N, M, k=k)) == 1
top_left = lambda N, M, k: np.flipud(np.tri(N, M, k=k)) == 1
bottom_left = lambda N, M, k: np.tri(N, M, k=k) == 1
top_right = lambda N, M, k: np.flipud(np.fliplr(np.tri(N, M, k=k))) == 1

def find_connected_nodes(image, i, nodes):
    out = []
    x, y = nodes[i]
    for j in range(i + 1, len(nodes)):
        if i == j: continue
        distance = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[j]), 2)))
        dst_x, dst_y = nodes[j]

        angle = np.arctan2(dst_y - y, dst_x - x)

        min_x = min(x, dst_x)
        min_y = min(y, dst_y)
        max_x = max(x, dst_x)
        max_y = max(y, dst_y)
        border = 2
        cutout = np.copy(image[min_y - border : max_y + border + 1, min_x - border : max_x + border + 1])

        if ((x - min_x) == 0 and (y - min_y) == 0) or ((dst_x - min_x) == 0 and (dst_y - min_y) == 0):
            cutout[top_right(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
            cutout[bottom_left(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
        else:
            cutout[top_left(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255
            cutout[bottom_right(cutout.shape[0], cutout.shape[1], min(cutout.shape) // 2 - cutout.shape[0])] = 255

        contours = find_tree_contours(cutout, 2)
        for contour in contours:
            other = False
            for k in range(len(nodes)):
                if k == j or k == i: continue
                k_x, k_y = nodes[k]
                k_x = k_x - min_x + border
                k_y = k_y - min_y + border
                if k_x >= border and k_x < cutout.shape[-1] - border and k_y >= border and k_y < cutout.shape[-2] - border:
                    if contour[k_y, k_x] < 200:
                        other = True
                        break

            # Check if there's not something along the way
            if False: #i == 43 and j == 57:
                print(i, j, other,
                      contour[y     - min_y + border,     x - min_x + border],
                      contour[dst_y - min_y + border, dst_x - min_x + border],
                      nodes[i], nodes[j])
                print(image[dst_y - 2 : dst_y + 3, dst_x - 2 : dst_x + 3])
                print(cutout[:5, :5])
                print(contour[:5, :5])

                cv2.imshow('1', cutout)
                cv2.imshow('2', contour)
                cv2.waitKey(0)

            if other == False and contour[y - min_y + border, x - min_x + border] < 200 and contour[dst_y - min_y + border, dst_x - min_x + border] < 200:
                out.append(j)
    return out

def find_nearest_connected_nodes(image, i, nodes):
    out = []
    x, y = nodes[i]
    min_j = i
    min_dist = np.max(image.shape)
    distances = []
    good_js = []
    for j in range(len(nodes)):
        distance = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[j]), 2)))
        distances.append(distance)
        if i == j: continue
        dst_x, dst_y = nodes[j]

        angle = np.arctan2(dst_y - y, dst_x - x)

        min_x = min(x, dst_x)
        min_y = min(y, dst_y)
        max_x = max(x, dst_x)
        max_y = max(y, dst_y)
        cutout = np.copy(image[min_y - 1  : max_y + 2, min_x - 1 : max_x + 2])

        contours = find_tree_contours(cutout, 1)
        for contour in contours:
            if contour[y - min_y + 1, x - min_x + 1] < 128 and contour[dst_y - min_y + 1, dst_x - min_x + 1] < 128:
                if distance < min_dist:
                    min_j = j
                    min_dist = distance
                good_js.append(j)
    for j in good_js:
        if distances[j] < min_dist * 1.5:
            out.append(nodes[j])
    return out

# Nodes from corners
def nodes_from_corners(image, gray, points, max_dist=16, iterations=1):
    if iterations == 0:
        out = []

        for j, point in enumerate(points):
            x, y = map(int, point)
            A = gray[y - 1 : y + 2, x - 1 : x + 2]
            off_y, off_x = np.array(np.unravel_index(A.argmin(), A.shape)) - 1
            new_x, new_y = x + off_x, y + off_y
            if gray[new_y, new_x] < gray[y, x]:
                out.append((new_x,new_y))
            else:
                out.append((x,y))

        return out
    print('iternations:', iterations)
    nodes = []
    while len(points):
        point = points[-1]
        points.pop()
        neighbors = [point]
        for i in reversed(range(len(points))):
            dist = np.sqrt(np.sum(np.power(np.array(point) - np.array(points[i]), 2)))
            if dist < max_dist:
                neighbors.append(points[i])
                points.pop(i)
        neighbors = np.array(neighbors)
        point = neighbors[np.random.randint(len(neighbors))]
        point = np.mean(neighbors, axis=0)
        print('Remaining corner points to assign a corner:', len(points))
        if iterations == 1:
            #image[point[1], point[0], :] = (255, 0, 0)
            cv2.circle(image, tuple(map(int, point)), 4, (255,128,0), 2)
        nodes.append(point)
    return nodes_from_corners(image, gray, nodes, max_dist + 2, iterations - 1)

# Find the contour of the tree
def find_tree_contour(gray, param, min_frac=0.5):
    _,thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV) # threshold
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    dilated = cv2.dilate(thresh,kernel,iterations = param) # dilate
    im2, contours, hierarchy = cv2.findContours(dilated, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    out = np.zeros(gray.shape, dtype=np.uint8) + 255

    for i, contour in enumerate(contours):
        [x,y,w,h] = cv2.boundingRect(contour)
        if gray.shape[-1] > gray.shape[-2]:
            if w < gray.shape[-1] * min_frac:
                continue
        elif gray.shape[-1] < gray.shape[-2]:
            if h < gray.shape[-2] * min_frac:
                continue
        else:
            if h < gray.shape[-2] * min_frac or w < gray.shape[-1] * min_frac:
                continue
        #cv2.rectangle(img, (x,y), (x+w,y+h), (255,0,255), 1)
        cv2.drawContours(out, contours, i, 0, -1)

    out2 = np.zeros(gray.shape, dtype=np.uint8) + 255
    out2[out == 0] = gray[out == 0]
    return out2

def find_tree_contours(gray, param):
    _,thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV) # threshold
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS,(3,3))
    dilated = cv2.dilate(thresh,kernel,iterations = param) # dilate
    im2, contours, hierarchy = cv2.findContours(dilated, cv2.RETR_LIST, cv2.CHAIN_APPROX_TC89_KCOS)

    outs = []
    for i, contour in enumerate(contours):
        [x,y,w,h] = cv2.boundingRect(contour)
        if h < gray.shape[0] * 0.5 or w < gray.shape[1] * 0.5: continue
        #cv2.rectangle(img, (x,y), (x+w,y+h), (255,0,255), 1)
        out = np.zeros(gray.shape, dtype=np.uint8) + 255
        cv2.drawContours(out, contours, i, 0, -1)
        outs.append(out)
    #out2[out < 2010] = 0
    return outs

def main():
    filename = sys.argv[1]
    param = int(sys.argv[2])

    img = cv2.imread(filename)
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
            img[y,x,:] = np.min(img[y,x])
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

    tree_contour = find_tree_contour(gray, param)
    cv2.imshow('dst2', tree_contour)

    dst = cv2.cornerHarris(tree_contour, 3, 3, 0.01)
    dst = cv2.dilate(dst, None)
    corners = dst > 0.05 * dst.max()

    points = []
    for y in range(img.shape[0]):
        for x in range(img.shape[1]):
            if corners[y,x]:
                points.append((x,y))

    nodes = nodes_from_corners(img, gray, points, max_dist=8, iterations=3)
    for i, node in enumerate(nodes):
        x, y = node
        print(i, gray[y, x])

    # Figure out node neighbors
    node_neighbors = {}
    for i, node in enumerate(nodes):
        node_neighbors[i] = find_connected_nodes(gray, i, nodes)

    # Do reciprocal adding of nodes
    for i in range(len(nodes)):
        for j in node_neighbors[i]:
            if i not in node_neighbors[j]:
                node_neighbors[j].append(i)

    # Find Triangles
    for i in range(len(nodes)):
        for j in range(i + 1, len(nodes)):
            for k in range(j + 1, len(nodes)):
                if (i in node_neighbors[j] and i in node_neighbors[k] and
                    j in node_neighbors[i] and j in node_neighbors[k] and
                    k in node_neighbors[i] and k in node_neighbors[j]):
                    d1 = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[j]), 2)))
                    d2 = np.sqrt(np.sum(np.power(np.array(nodes[i]) - np.array(nodes[k]), 2)))
                    d3 = np.sqrt(np.sum(np.power(np.array(nodes[j]) - np.array(nodes[k]), 2)))
                    max_d = max(d1, d2, d3)
                    if d1 == max_d:
                        node_neighbors[i] = [x for x in node_neighbors[i] if x != j]
                        node_neighbors[j] = [x for x in node_neighbors[j] if x != i]
                    elif d2 == max_d:
                        node_neighbors[i] = [x for x in node_neighbors[i] if x != k]
                        node_neighbors[k] = [x for x in node_neighbors[k] if x != i]
                    else:
                        node_neighbors[j] = [x for x in node_neighbors[j] if x != k]
                        node_neighbors[k] = [x for x in node_neighbors[k] if x != j]

                    print('Triangle', i, j, k)

    # Delete nodes with only 2 neighbors as they are connections:
    for i in range(len(nodes)):
        if len(node_neighbors[i]) == 2:
            n1, n2 = node_neighbors[i]
            continue
            for j in node_neighbors:
                print(node_neighbors[j], i, j)
                node_neighbors[j] = [x for x in node_neighbors[j] if x != i]
            node_neighbors[n1].append(n2)
            node_neighbors[n2].append(n1)
            del node_neighbors[i]

    # Print lines
    for i in sorted(node_neighbors):
        node_x, node_y = nodes[i]
        for neighbor in node_neighbors[i]:
            neighbor_x, neighbor_y = nodes[neighbor]
            cv2.line(img, (node_x, node_y), (neighbor_x, neighbor_y), (255,0,0), 2)

    for i in sorted(node_neighbors):
        print(i, node_neighbors[i])

    for i in sorted(node_neighbors):
        cv2.putText(img, str(i), nodes[i], cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))

    cv2.imshow('dst', img)
    if cv2.waitKey(0) & 0xff == 27:
        cv2.destroyAllWindows()

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


<commit_before>from hydromet import disaggregate
from hydromet import models
from hydromet import stats

from ._version import get_versions
__version__ = get_versions()['version']
del get_versions
<commit_msg>Disable import so that hydromet can be used w/o hydromath.so<commit_after>from hydromet import disaggregate
#from hydromet import models
from hydromet import stats

from ._version import get_versions
__version__ = get_versions()['version']
del get_versions
<|endoftext|>


<commit_before>from collections import defaultdict
import os
import random
import sqlite3
import uuid

# FIXME these are hardcoded in some SQL statements below.  SQLite does not
# seem to have the concept of variables...
INITIALIZED = 0
ASSIGNED = 1
SUCCESSFUL = 2
FAILED = 3
ABORTED = 4
INCOMPLETE = 5
PUBLISHED = 6

class JobitStore:
    def __init__(self, config):
        self.uuid = str(uuid.uuid4()).replace('-', '')
        self.config = config
        self.db_path = os.path.join(config['workdir'], "lobster.db")
        self.db = sqlite3.connect(self.db_path)

        # Use four databases: one for jobits, jobs, hosts, datasets each
        self.db.execute("""create table if not exists datasets(
            id integer primary key autoincrement,
            dataset text,
            label text,
            path text,
            release text,
            global_tag text,
            publish_label text,
            pset_hash text default null,
            cfg text,
            uuid text,
            jobsize text,
            file_based int,
            jobits integer,
            jobits_running int default 0,
            jobits_done int default 0,
            jobits_left int default 0,
            events int default 0,
            events_read int default 0,
            events_written int default 0,
            bytes_input int default 0,
            bytes_output int default 0)""")
        self.db.execute("""create table if not exists jobs(
            id integer primary key autoincrement,
            host text,
            dataset int,
            published_file_block text,
            status int default 0,
            exit_code int,
            submissions int default 0,
            jobits int default 0,
            jobits_processed int default 0,
            events_read int default 0,
            events_written int default 0,
            time_submit int,
            time_transfer_in_start int,
            time_transfer_in_end int,
            time_wrapper_start int,
            time_wrapper_ready int,
            time_file_requested int,
            time_file_opened int,
            time_file_processing int,
            time_processing_end int,
            time_chirp_end int,
            time_transfer_out_start int,
            time_transfer_out_end int,
            time_retrieved int,
            time_on_worker int,
            time_total_on_worker int,
            time_cpu int,
            bytes_received int,
            bytes_sent int,
            bytes_output int default 0,
            foreign key(dataset) references datasets(id))""")
        self.db.commit()

        try:
            cur = self.db.execute("select max(id) from jobs")
            count = int(cur.fetchone()[0])
        except:
            pass

    def disconnect(self):
        self.db.close()

    def register(self, dataset_cfg, dataset_info):
        label = dataset_cfg['label']
        unique_args = dataset_cfg.get('unique parameters', [None])

        cur = self.db.cursor()
        cur.execute("""insert into datasets
                       (dataset,
                       label,
                       path,
                       release,
                       global_tag,
                       publish_label,
                       cfg,
                       uuid,
                       file_based,
                       jobsize,
                       jobits,
                       jobits_left,
                       events,
                       bytes_input)
                       values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", (
                           dataset_cfg.get('dataset', dataset_cfg.get('files', None)),
                           label,
                           os.path.join(self.config['stageout location'], label),
                           os.path.basename(os.environ['LOCALRT']),
                           dataset_cfg.get('global tag'),
                           dataset_cfg.get('publish label', dataset_cfg['label']).replace('-', '_'), #TODO: more lexical checks #TODO: publish label check
                           dataset_cfg.get('cmssw config'),
                           self.uuid,
                           dataset_info.file_based,
                           dataset_info.jobsize,
                           dataset_info.total_lumis * len(unique_args),
                           dataset_info.total_lumis * len(unique_args),
                           dataset_info.total_events,
                           sum(dataset_info.filesizes.values())))
        dset_id = cur.lastrowid

        self.db.execute("""create table if not exists files_{0}(
            id integer primary key autoincrement,
            filename text,
            skipped int default 0,
            jobits int,
            jobits_done int default 0,
            jobits_running int default 0,
            events int,
            events_read int default 0,
            bytes int default 0)""".format(label))

        cur.execute("""create table if not exists jobits_{0}(
            id integer primary key autoincrement,
            job integer,
            run integer,
            lumi integer,
            file integer,
            status integer default 0,
            arg text,
            foreign key(job) references jobs(id),
            foreign key(file) references files(id))""".format(label))

        for file in dataset_info.files:
            file_lumis = len(dataset_info.lumis[file])
            cur.execute(
                    """insert into files_{0}(jobits, events, filename, bytes) values (?, ?, ?, ?)""".format(label), (
                        file_lumis * len(unique_args),
                        dataset_info.event_counts[file],
                        file,
                        dataset_info.filesizes[file]))
            file_id = cur.lastrowid

            for arg in unique_args:
                columns = [(file_id, run, lumi, arg) for (run, lumi) in dataset_info.lumis[file]]
                self.db.executemany("insert into jobits_{0}(file, run, lumi, arg) values (?, ?, ?, ?)".format(label), columns)

        self.db.execute("create index if not exists index_filename_{0} on files_{0}(filename)".format(label))
        self.db.execute("create index if not exists index_events_{0} on jobits_{0}(run, lumi)".format(label))
        self.db.execute("create index if not exists index_files_{0} on jobits_{0}(file)".format(label))

        self.db.commit()

    def pop_jobits(self, num=1, bijective=False):
        """
        Create a predetermined number of jobs.  The task these are
        created for is drawn randomly from all unfinished tasks.

        Arguments:
            num: the number of jobs to be created (default 1)
            bijective: process lumis from one file only, where possible (default False)
        Returns:
            a list containing an id, dataset label, file information (id,
            filename), lumi information (id, file id, run, lumi)
        """

        rows = [xs for xs in self.db.execute("""
            select label, id, jobits - jobits_done - jobits_running, jobsize
            from datasets
            where jobits_done + jobits_running < jobits""")]
        if len(rows) == 0:
            return None

        dataset, dataset_id, remaining, jobsize = random.choice(rows)
        size = [int(jobsize)] * num

        fileinfo = list(self.db.execute("""select id, filename
                    from files_{0}
                    where jobits_done + jobits_running < jobits
                    order by skipped asc""".format(dataset)))
        files = [x for (x, y) in fileinfo]
        fileinfo = dict(fileinfo)

        rows = []
        if bijective:
            remaining = sum(size)
            size = []
            for file in files:
                lumis = self.db.execute("""
                    select id, file, run, lumi, arg
                    from jobits_{0}
                    where file=? and (status<>1 and status<>2 and status<>6)""".format(dataset), (file,))
                size.append(len(lumis))
                remaining -= len(lumis)
                if remaining <= 0:
                    break
                rows += lumis
        else:
            for i in range(0, len(files), 40):
                chunk = files[i:i + 40]
                rows.extend(self.db.execute("""
                    select id, file, run, lumi, arg
                    from jobits_{0}
                    where file in ({1}) and (status<>1 and status<>2 and status<>6)
                    """.format(dataset, ', '.join('?' for _ in chunk)), chunk))

        # files and lumis for individual jobs
        files = set()
        jobits = []

        # job container and current job size
        jobs = []
        current_size = 0

        for id, file, run, lumi, arg in rows:
            if current_size == 0:
                if len(size) == 0:
                    break
                cur = self.db.cursor()
                cur.execute("insert into jobs(dataset, status) values (?, 1)", (dataset_id,))
                job_id = cur.lastrowid

            if lumi > 0:
                for (ls_id, ls_file, ls_run, ls_lumi) in self.db.execute("""
                        select id, file, run, lumi
                        from jobits_{0}
                        where run=? and lumi=? and status not in (1, 2, 6)""".format(dataset), (run, lumi)):
                    jobits.append((ls_id, ls_file, ls_run, ls_lumi))
                    files.add(ls_file)
            else:
                jobits.append((id, file, run, lumi))
                files.add(file)

            current_size += 1

            if current_size == size[0]:
                jobs.append((
                    str(job_id),
                    dataset,
                    [(id, fileinfo[id]) for id in files],
                    jobits,
                    arg))

                files = set()
                jobits = []

                current_size = 0
                size.pop(0)

        if current_size > 0:
            jobs.append((
                str(job_id),
                dataset,
                [(id, fileinfo[id]) for id in files],
                jobits,
                arg))

        dataset_update = []
        file_update = defaultdict(int)
        job_update = defaultdict(int)
        jobit_update = []

        for (job, label, files, jobits, arg) in jobs:
            dataset_update += jobits
            job_update[job] = len(jobits)
            jobit_update += [(job, id) for (id, file, run, lumi) in jobits]
            for (id, filename) in files:
                file_update[id] += len(filter(lambda tpl: tpl[1] == id, jobits))

        self.db.execute(
                "update datasets set jobits_running=(jobits_running + ?) where id=?",
                (len(dataset_update), dataset_id))

        self.db.executemany("update files_{0} set jobits_running=(jobits_running + ?) where id=?".format(dataset),
                [(v, k) for (k, v) in file_update.items()])
        self.db.executemany("update jobs set jobits=? where id=?",
                [(v, k) for (k, v) in job_update.items()])
        self.db.executemany("update jobits_{0} set status=1, job=? where id=?".format(dataset),
                jobit_update)

        self.db.commit()

        return jobs if len(jobit_update) > 0 else None

    def reset_jobits(self):
        with self.db as db:
            ids = [id for (id,) in db.execute("select id from jobs where status=1")]
            db.execute("update datasets set jobits_running=0")
            db.execute("update jobs set status=4 where status=1")
            for (label,) in db.execute("select label from datasets"):
                db.execute("update files_{0} set jobits_running=0".format(label))
                db.execute("update jobits_{0} set status=4 where status=1".format(label))
        return ids

    def update_jobits(self, jobinfos):
        job_updates = []
        dset_infos = {}

        for (dset, updates) in jobinfos.items():
            file_updates = []
            jobit_updates = []
            jobit_generic_updates = []

            for (job_update, file_update, jobit_update) in updates:
                job_updates.append(job_update)
                file_updates += file_update

                # jobits either fail or are successful
                jobit_status = FAILED if job_update[-2] == FAILED else SUCCESSFUL

                jobit_updates += jobit_update
                # the last entry in the job_update is the id
                jobit_generic_updates.append((jobit_status, job_update[-1]))

                # bytes written, events read and written
                try:
                    dset_infos[dset][0] += job_update[-6]
                    dset_infos[dset][1] += job_update[-4]
                    dset_infos[dset][2] += job_update[-3]
                except KeyError:
                    dset_infos[dset] = [job_update[-6], job_update[-4], job_update[-3]]

            # update all jobits of the jobs
            self.db.executemany("""update jobits_{0} set
                status=?
                where job=?""".format(dset),
                jobit_generic_updates)

            # update selected, missed jobits
            self.db.executemany("""update jobits_{0} set
                status=?
                where id=?""".format(dset),
                jobit_updates)

            # update files in the dataset
            self.db.executemany("""update files_{0} set
                jobits_running=(select count(*) from jobits_{0} where status==1 and file=files_{0}.id),
                jobits_done=(select count(*) from jobits_{0} where status==2 and file=files_{0}.id),
                events_read=(events_read + ?),
                skipped=(skipped + ?)
                where id=?""".format(dset),
                file_updates)

        self.db.executemany("""update jobs set
            host=?,
            exit_code=?,
            submissions=?,
            time_submit=?,
            time_transfer_in_start=?,
            time_transfer_in_end=?,
            time_wrapper_start=?,
            time_wrapper_ready=?,
            time_file_requested=?,
            time_file_opened=?,
            time_file_processing=?,
            time_processing_end=?,
            time_chirp_end=?,
            time_transfer_out_start=?,
            time_transfer_out_end=?,
            time_retrieved=?,
            time_on_worker=?,
            time_total_on_worker=?,
            time_cpu=?,
            bytes_received=?,
            bytes_sent=?,
            bytes_output=?,
            jobits_processed=(jobits - ?),
            events_read=?,
            events_written=?,
            status=?
            where id=?""",
            job_updates)

        for label in jobinfos.keys():
            self.update_dataset_stats(label, *dset_infos[label])

        self.db.commit()

    def update_dataset_stats(self, label, bytes_written=0, events_read=0, events_written=0):
        file_based = self.db.execute("select file_based from datasets where label=?", (label,)).fetchone()[0]

        self.db.execute("""
            update datasets set
                jobits_running=(select count(*) from jobits_{0} where status==1),
                jobits_done=(select count(*) from jobits_{0} where status==2),
                jobits_left=(select count(*) from jobits_{0} where status not in (1, 2)),
                events_read=(events_read + ?),
                events_written=(events_written + ?),
                bytes_output=(bytes_output + ?)
            where label=?""".format(label),
            (events_read, events_written, bytes_written, label))

    def unfinished_jobits(self):
        cur = self.db.execute("select sum(jobits - jobits_done) from datasets")
        return cur.fetchone()[0]

    def running_jobits(self):
        cur = self.db.execute("select sum(jobits_running) from datasets")
        return cur.fetchone()[0]

    def dataset_info(self, label):
        cur = self.db.execute("""select dataset,
            path,
            release,
            global_tag,
            publish_label,
            cfg,
            pset_hash,
            id,
            uuid
            from datasets
            where label==?""", (label,))

        return cur.fetchone()

    def finished_jobs(self, dataset):
        cur = self.db.execute("""select jobs.id
            from jobs, datasets
            where jobs.status=?
            and datasets.label=?
            and jobs.dataset==datasets.id""", (SUCCESSFUL, dataset))

        return cur.fetchall()

    def update_published(self, blocks):
        columns = [(PUBLISHED, block, id) for block, id in blocks]

        self.db.executemany("""update jobs
            set status=?,
            published_file_block=?
            where id=?""", columns)

        self.db.executemany("""update jobits
            set status=?
            where id=?""", [(x, z) for x, y, z in columns])

        self.db.commit()

    def update_datasets(self, column, value, label):
        self.db.execute("""update datasets
            set %s=?
            where label=?""" % column, (value, label,))

        self.db.commit()
<commit_msg>Re-add split lumi section double-processing veto.<commit_after>from collections import defaultdict
import os
import random
import sqlite3
import uuid

# FIXME these are hardcoded in some SQL statements below.  SQLite does not
# seem to have the concept of variables...
INITIALIZED = 0
ASSIGNED = 1
SUCCESSFUL = 2
FAILED = 3
ABORTED = 4
INCOMPLETE = 5
PUBLISHED = 6

class JobitStore:
    def __init__(self, config):
        self.uuid = str(uuid.uuid4()).replace('-', '')
        self.config = config
        self.db_path = os.path.join(config['workdir'], "lobster.db")
        self.db = sqlite3.connect(self.db_path)

        # Use four databases: one for jobits, jobs, hosts, datasets each
        self.db.execute("""create table if not exists datasets(
            id integer primary key autoincrement,
            dataset text,
            label text,
            path text,
            release text,
            global_tag text,
            publish_label text,
            pset_hash text default null,
            cfg text,
            uuid text,
            jobsize text,
            file_based int,
            jobits integer,
            jobits_running int default 0,
            jobits_done int default 0,
            jobits_left int default 0,
            events int default 0,
            events_read int default 0,
            events_written int default 0,
            bytes_input int default 0,
            bytes_output int default 0)""")
        self.db.execute("""create table if not exists jobs(
            id integer primary key autoincrement,
            host text,
            dataset int,
            published_file_block text,
            status int default 0,
            exit_code int,
            submissions int default 0,
            jobits int default 0,
            jobits_processed int default 0,
            events_read int default 0,
            events_written int default 0,
            time_submit int,
            time_transfer_in_start int,
            time_transfer_in_end int,
            time_wrapper_start int,
            time_wrapper_ready int,
            time_file_requested int,
            time_file_opened int,
            time_file_processing int,
            time_processing_end int,
            time_chirp_end int,
            time_transfer_out_start int,
            time_transfer_out_end int,
            time_retrieved int,
            time_on_worker int,
            time_total_on_worker int,
            time_cpu int,
            bytes_received int,
            bytes_sent int,
            bytes_output int default 0,
            foreign key(dataset) references datasets(id))""")
        self.db.commit()

        try:
            cur = self.db.execute("select max(id) from jobs")
            count = int(cur.fetchone()[0])
        except:
            pass

    def disconnect(self):
        self.db.close()

    def register(self, dataset_cfg, dataset_info):
        label = dataset_cfg['label']
        unique_args = dataset_cfg.get('unique parameters', [None])

        cur = self.db.cursor()
        cur.execute("""insert into datasets
                       (dataset,
                       label,
                       path,
                       release,
                       global_tag,
                       publish_label,
                       cfg,
                       uuid,
                       file_based,
                       jobsize,
                       jobits,
                       jobits_left,
                       events,
                       bytes_input)
                       values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", (
                           dataset_cfg.get('dataset', dataset_cfg.get('files', None)),
                           label,
                           os.path.join(self.config['stageout location'], label),
                           os.path.basename(os.environ['LOCALRT']),
                           dataset_cfg.get('global tag'),
                           dataset_cfg.get('publish label', dataset_cfg['label']).replace('-', '_'), #TODO: more lexical checks #TODO: publish label check
                           dataset_cfg.get('cmssw config'),
                           self.uuid,
                           dataset_info.file_based,
                           dataset_info.jobsize,
                           dataset_info.total_lumis * len(unique_args),
                           dataset_info.total_lumis * len(unique_args),
                           dataset_info.total_events,
                           sum(dataset_info.filesizes.values())))
        dset_id = cur.lastrowid

        self.db.execute("""create table if not exists files_{0}(
            id integer primary key autoincrement,
            filename text,
            skipped int default 0,
            jobits int,
            jobits_done int default 0,
            jobits_running int default 0,
            events int,
            events_read int default 0,
            bytes int default 0)""".format(label))

        cur.execute("""create table if not exists jobits_{0}(
            id integer primary key autoincrement,
            job integer,
            run integer,
            lumi integer,
            file integer,
            status integer default 0,
            arg text,
            foreign key(job) references jobs(id),
            foreign key(file) references files(id))""".format(label))

        for file in dataset_info.files:
            file_lumis = len(dataset_info.lumis[file])
            cur.execute(
                    """insert into files_{0}(jobits, events, filename, bytes) values (?, ?, ?, ?)""".format(label), (
                        file_lumis * len(unique_args),
                        dataset_info.event_counts[file],
                        file,
                        dataset_info.filesizes[file]))
            file_id = cur.lastrowid

            for arg in unique_args:
                columns = [(file_id, run, lumi, arg) for (run, lumi) in dataset_info.lumis[file]]
                self.db.executemany("insert into jobits_{0}(file, run, lumi, arg) values (?, ?, ?, ?)".format(label), columns)

        self.db.execute("create index if not exists index_filename_{0} on files_{0}(filename)".format(label))
        self.db.execute("create index if not exists index_events_{0} on jobits_{0}(run, lumi)".format(label))
        self.db.execute("create index if not exists index_files_{0} on jobits_{0}(file)".format(label))

        self.db.commit()

    def pop_jobits(self, num=1, bijective=False):
        """
        Create a predetermined number of jobs.  The task these are
        created for is drawn randomly from all unfinished tasks.

        Arguments:
            num: the number of jobs to be created (default 1)
            bijective: process lumis from one file only, where possible (default False)
        Returns:
            a list containing an id, dataset label, file information (id,
            filename), lumi information (id, file id, run, lumi)
        """

        rows = [xs for xs in self.db.execute("""
            select label, id, jobits - jobits_done - jobits_running, jobsize
            from datasets
            where jobits_done + jobits_running < jobits""")]
        if len(rows) == 0:
            return None

        dataset, dataset_id, remaining, jobsize = random.choice(rows)
        size = [int(jobsize)] * num

        fileinfo = list(self.db.execute("""select id, filename
                    from files_{0}
                    where jobits_done + jobits_running < jobits
                    order by skipped asc""".format(dataset)))
        files = [x for (x, y) in fileinfo]
        fileinfo = dict(fileinfo)

        rows = []
        if bijective:
            remaining = sum(size)
            size = []
            for file in files:
                lumis = self.db.execute("""
                    select id, file, run, lumi, arg
                    from jobits_{0}
                    where file=? and (status<>1 and status<>2 and status<>6)""".format(dataset), (file,))
                size.append(len(lumis))
                remaining -= len(lumis)
                if remaining <= 0:
                    break
                rows += lumis
        else:
            for i in range(0, len(files), 40):
                chunk = files[i:i + 40]
                rows.extend(self.db.execute("""
                    select id, file, run, lumi, arg
                    from jobits_{0}
                    where file in ({1}) and (status<>1 and status<>2 and status<>6)
                    """.format(dataset, ', '.join('?' for _ in chunk)), chunk))

        # files and lumis for individual jobs
        files = set()
        jobits = []

        # lumi veto to avoid duplicated processing
        all_lumis = set()

        # job container and current job size
        jobs = []
        current_size = 0

        for id, file, run, lumi, arg in rows:
            if (run, lumi) in all_lumis:
                continue

            if current_size == 0:
                if len(size) == 0:
                    break
                cur = self.db.cursor()
                cur.execute("insert into jobs(dataset, status) values (?, 1)", (dataset_id,))
                job_id = cur.lastrowid

            if lumi > 0:
                all_lumis.add((run, lumi))
                for (ls_id, ls_file, ls_run, ls_lumi) in self.db.execute("""
                        select id, file, run, lumi
                        from jobits_{0}
                        where run=? and lumi=? and status not in (1, 2, 6)""".format(dataset), (run, lumi)):
                    jobits.append((ls_id, ls_file, ls_run, ls_lumi))
                    files.add(ls_file)
            else:
                jobits.append((id, file, run, lumi))
                files.add(file)

            current_size += 1

            if current_size == size[0]:
                jobs.append((
                    str(job_id),
                    dataset,
                    [(id, fileinfo[id]) for id in files],
                    jobits,
                    arg))

                files = set()
                jobits = []

                current_size = 0
                size.pop(0)

        if current_size > 0:
            jobs.append((
                str(job_id),
                dataset,
                [(id, fileinfo[id]) for id in files],
                jobits,
                arg))

        dataset_update = []
        file_update = defaultdict(int)
        job_update = defaultdict(int)
        jobit_update = []

        for (job, label, files, jobits, arg) in jobs:
            dataset_update += jobits
            job_update[job] = len(jobits)
            jobit_update += [(job, id) for (id, file, run, lumi) in jobits]
            for (id, filename) in files:
                file_update[id] += len(filter(lambda tpl: tpl[1] == id, jobits))

        self.db.execute(
                "update datasets set jobits_running=(jobits_running + ?) where id=?",
                (len(dataset_update), dataset_id))

        self.db.executemany("update files_{0} set jobits_running=(jobits_running + ?) where id=?".format(dataset),
                [(v, k) for (k, v) in file_update.items()])
        self.db.executemany("update jobs set jobits=? where id=?",
                [(v, k) for (k, v) in job_update.items()])
        self.db.executemany("update jobits_{0} set status=1, job=? where id=?".format(dataset),
                jobit_update)

        self.db.commit()

        return jobs if len(jobit_update) > 0 else None

    def reset_jobits(self):
        with self.db as db:
            ids = [id for (id,) in db.execute("select id from jobs where status=1")]
            db.execute("update datasets set jobits_running=0")
            db.execute("update jobs set status=4 where status=1")
            for (label,) in db.execute("select label from datasets"):
                db.execute("update files_{0} set jobits_running=0".format(label))
                db.execute("update jobits_{0} set status=4 where status=1".format(label))
        return ids

    def update_jobits(self, jobinfos):
        job_updates = []
        dset_infos = {}

        for (dset, updates) in jobinfos.items():
            file_updates = []
            jobit_updates = []
            jobit_generic_updates = []

            for (job_update, file_update, jobit_update) in updates:
                job_updates.append(job_update)
                file_updates += file_update

                # jobits either fail or are successful
                jobit_status = FAILED if job_update[-2] == FAILED else SUCCESSFUL

                jobit_updates += jobit_update
                # the last entry in the job_update is the id
                jobit_generic_updates.append((jobit_status, job_update[-1]))

                # bytes written, events read and written
                try:
                    dset_infos[dset][0] += job_update[-6]
                    dset_infos[dset][1] += job_update[-4]
                    dset_infos[dset][2] += job_update[-3]
                except KeyError:
                    dset_infos[dset] = [job_update[-6], job_update[-4], job_update[-3]]

            # update all jobits of the jobs
            self.db.executemany("""update jobits_{0} set
                status=?
                where job=?""".format(dset),
                jobit_generic_updates)

            # update selected, missed jobits
            self.db.executemany("""update jobits_{0} set
                status=?
                where id=?""".format(dset),
                jobit_updates)

            # update files in the dataset
            self.db.executemany("""update files_{0} set
                jobits_running=(select count(*) from jobits_{0} where status==1 and file=files_{0}.id),
                jobits_done=(select count(*) from jobits_{0} where status==2 and file=files_{0}.id),
                events_read=(events_read + ?),
                skipped=(skipped + ?)
                where id=?""".format(dset),
                file_updates)

        self.db.executemany("""update jobs set
            host=?,
            exit_code=?,
            submissions=?,
            time_submit=?,
            time_transfer_in_start=?,
            time_transfer_in_end=?,
            time_wrapper_start=?,
            time_wrapper_ready=?,
            time_file_requested=?,
            time_file_opened=?,
            time_file_processing=?,
            time_processing_end=?,
            time_chirp_end=?,
            time_transfer_out_start=?,
            time_transfer_out_end=?,
            time_retrieved=?,
            time_on_worker=?,
            time_total_on_worker=?,
            time_cpu=?,
            bytes_received=?,
            bytes_sent=?,
            bytes_output=?,
            jobits_processed=(jobits - ?),
            events_read=?,
            events_written=?,
            status=?
            where id=?""",
            job_updates)

        for label in jobinfos.keys():
            self.update_dataset_stats(label, *dset_infos[label])

        self.db.commit()

    def update_dataset_stats(self, label, bytes_written=0, events_read=0, events_written=0):
        file_based = self.db.execute("select file_based from datasets where label=?", (label,)).fetchone()[0]

        self.db.execute("""
            update datasets set
                jobits_running=(select count(*) from jobits_{0} where status==1),
                jobits_done=(select count(*) from jobits_{0} where status==2),
                jobits_left=(select count(*) from jobits_{0} where status not in (1, 2)),
                events_read=(events_read + ?),
                events_written=(events_written + ?),
                bytes_output=(bytes_output + ?)
            where label=?""".format(label),
            (events_read, events_written, bytes_written, label))

    def unfinished_jobits(self):
        cur = self.db.execute("select sum(jobits - jobits_done) from datasets")
        return cur.fetchone()[0]

    def running_jobits(self):
        cur = self.db.execute("select sum(jobits_running) from datasets")
        return cur.fetchone()[0]

    def dataset_info(self, label):
        cur = self.db.execute("""select dataset,
            path,
            release,
            global_tag,
            publish_label,
            cfg,
            pset_hash,
            id,
            uuid
            from datasets
            where label==?""", (label,))

        return cur.fetchone()

    def finished_jobs(self, dataset):
        cur = self.db.execute("""select jobs.id
            from jobs, datasets
            where jobs.status=?
            and datasets.label=?
            and jobs.dataset==datasets.id""", (SUCCESSFUL, dataset))

        return cur.fetchall()

    def update_published(self, blocks):
        columns = [(PUBLISHED, block, id) for block, id in blocks]

        self.db.executemany("""update jobs
            set status=?,
            published_file_block=?
            where id=?""", columns)

        self.db.executemany("""update jobits
            set status=?
            where id=?""", [(x, z) for x, y, z in columns])

        self.db.commit()

    def update_datasets(self, column, value, label):
        self.db.execute("""update datasets
            set %s=?
            where label=?""" % column, (value, label,))

        self.db.commit()
<|endoftext|>


<commit_before>#from opengmcore import _opengmcore.adder as adder
from opengmcore   import *
from __version__                    import version
from functionhelper                 import *
from _inf_param                     import _MetaInfParam , InfParam
from _visu                          import visualizeGm
from _misc                          import defaultAccumulator

from __version__ import version
import time

from _inference_interface_generator import _inject_interface , InferenceBase

import inference
import hdf5
import benchmark

# initialize solver/ inference dictionaries
_solverDicts=[
   (inference.adder.minimizer.solver.__dict__ ,     'adder',       'minimizer' ),
   (inference.adder.maximizer.solver.__dict__,      'adder',       'maximizer' ),
   (inference.multiplier.integrator.solver.__dict__,'adder',       'integrator'),
   (inference.multiplier.minimizer.solver.__dict__, 'multiplier',  'minimizer' ),
   (inference.multiplier.maximizer.solver.__dict__, 'multiplier',  'maximizer' ),
   (inference.multiplier.integrator.solver.__dict__,'multiplier',  'integrator')
]
for infClass,infName in _inject_interface(_solverDicts):
  inference.__dict__[infName]=infClass

class Timer(object):
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        if self.name:
            print '[%s]' % self.name
        self.tstart = time.time()

    def __exit__(self, type, value, traceback):
        #if self.name:
        #    print '[%s]' % self.name,
        print '   Elapsed: %s' % (time.time() - self.tstart)

def saveGm(gm,f,d='gm'):
  """ save a graphical model to a hdf5 file:
  Args:
    gm : graphical model to save
    f  : filepath
    g  : dataset (defaut : 'gm')
  """
  hdf5.saveGraphicalModel(f,d)

def loadGm(f,d='gm',operator='adder'):
  """ save a graphical model to a hdf5 file:
  Args:
    f  : filepath
    g  : dataset (defaut : 'gm')
    operator : operator of the graphical model ('adder' / 'multiplier')
  """
  if(operator=='adder'):
    gm=adder.GraphicalModel()
  elif(operator=='multiplier'):
    gm=multiplier.GraphicalModel()
  else:
    raise RuntimeError("unknown operator: "+ operator)
  hdf5.loadGraphicalModel(gm,f,d)
  return gm

class TestModels(object):
  @staticmethod
  def chain3(nVar,nLabels):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-2):
      f=numpy.random.rand(nLabels,nLabels,nLabels)
      model.addFactor(model.addFunction(f),[x0,x0+1,x0+2])
    return model

  @staticmethod
  def chain4(nVar,nLabels):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-3):
      f=numpy.random.rand(nLabels,nLabels,nLabels,nLabels)
      model.addFactor(model.addFunction(f),[x0,x0+1,x0+2,x0+3])
    return model

  @staticmethod
  def chainN(nVar,nLabels,order,nSpecialUnaries=0,beta=1.0):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)

    for sn in range(nSpecialUnaries):
      r=int(numpy.random.rand(1)*nVar-1)
      rl=int(numpy.random.rand(1)*nLabels-1)

      unaries[r,rl]=0.0

    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-(order-1)):
      f=numpy.random.rand( *([nLabels]*order))
      f*=beta
      vis=numpy.arange(order)
      vis+=x0

      model.addFactor(model.addFunction(f),vis)
    return model

  @staticmethod
  def secondOrderGrid(dx,dy,nLabels):
    nVar=dx*dy
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    vis2Order=secondOrderGridVis(dx,dy,True)

    nF2=len(vis2Order)#.shape[0]
    f2s=numpy.random.rand(nF2,nLabels)

    model.addFactors(model.addFunctions(f2s),vis2Order)

    return model

class GenericTimingVisitor(object):
    def __init__(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        self.visitNth=visitNth
        self.reserve=reserve
        self.verbose=verbose
        self.multiline=multiline

        self.values_     = None
        self.runtimes_   = None
        self.bounds_     = None
        self.iterations_ = None

        self.t0          = None
        self.t1          = None
        self.iterNr      = 0
    def getValues(self):
        return numpy.require(self.values_,dtype=value_type)
    def getTimes(self):
        return numpy.require(self.runtimes_,dtype=value_type)
    def getBounds(self):
        return numpy.require(self.bounds_,dtype=value_type)
    def getIterations(self):
        return numpy.require(self.iterations_,dtype=value_type)

    def begin(self,inf):
      v = inf.value()
      b = inf.bound()

      self.values_    =[v]
      self.bounds_    =[b]
      self.runtimes_  =[0.0]
      self.iterations_=[self.iterNr]
      if self.verbose :
        print 'Begin :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)

      # start the timing
      self.t0         =time.time()
      self.t1         =time.time()
    def visit(self,inf):

      if(self.iterNr==0 or self.iterNr%self.visitNth==0):

        # "stop the timing"
        self.t1=time.time()

        # get the runtime of the run
        rt=self.t1-self.t0
        v = inf.value()
        b = inf.bound()
        if self.verbose :
          print 'Step  :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)

        # store results
        self.values_.append(v)
        self.bounds_.append(b)
        self.runtimes_.append(rt)
        self.iterations_.append(self.iterNr)

        # increment iteration number
        self.iterNr+=1

        # restart the timing
        self.t0=time.time()

      else:
        # increment iteration number
        self.iterNr+=1

    def end(self,inf):
        # "stop the timing"
        self.t1=time.time()

        # get the runtime of the run
        rt=self.t1-self.t0
        v = inf.value()
        b = inf.bound()
        if self.verbose :
          print 'End :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)
        # store results
        self.values_.append(v)
        self.bounds_.append(b)
        self.runtimes_.append(rt)
        self.iterations_.append(self.iterNr)

class __RandomFusion__(object):
    def __init__(self,gm,accumulator=None,parameter=InfParam()):

        if accumulator is None:
            self.accumulator=defaultAccumulator(gm=gm)
        else:
            self.accumulator=accumulator
        kwargs=parameter.kwargs
        self.gm_=gm

        self.steps = kwargs.get('steps', 100)
        self.fusionSolver = kwargs.get('fuisionSolver', 'lf2')

        self.arg_  = None
        self.value_ = None

        self.fusionMover=inference.adder.minimizer.FusionMover(self.gm_)

        self.nLabels = self.gm_.numberOfLabels(0)
        self.nVar    = self.gm_.numberOfVariables

    def timingVisitor(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        return GenericTimingVisitor(visitNth,reserve,verbose,multiline)

    def setStartingPoint(self,arg):
      self.arg_=arg
      self.value_=gm.evaluate(self.arg_)

    def infer(self,visitor=None):
        if(self.arg_ is None):
          self.arg_ = numpy.zeros(self.gm_.numberOfVariables,dtype=label_type)
          self.value_ = self.value_=self.gm_.evaluate(self.arg_)

        # start inference
        if visitor is not None:
          visitor.begin(self)

        # start fusion moves
        for x in range(self.steps):
          randState=numpy.random.randint(low=0, high=self.nLabels, size=self.nVar).astype(label_type)
          r = self.fusionMover.fuse(self.arg_,randState,self.fusionSolver)
          self.arg_=r[0]
          self.value_=r[1]
          visitor.visit(self)

        # end inference
        if visitor is not None:
          visitor.end(self)

    def name(self):
      return "RandomFusion"

    def bound(self):
      return -1.0*float('inf')

    def arg(self):
        return self.arg_

    def value(self):
        return self.value_

class __CheapInitialization__(object):
    def __init__(self,gm,accumulator=None,parameter=InfParam()):

        if accumulator is None:
            self.accumulator=defaultAccumulator(gm=gm)
        else:
            self.accumulator=accumulator
        kwargs=parameter.kwargs
        self.gm_=gm

        self.arg_  = None
        self.value_ = None
        self.initType = kwargs.get('initType', 'localOpt')

    def timingVisitor(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        return GenericTimingVisitor(visitNth,reserve,verbose,multiline)

    def setStartingPoint(self,arg):
      self.arg_=arg
      self.value_=gm.evaluate(self.arg_)

    def infer(self,visitor=None):
        if(self.arg_ is None):
          self.arg_ = numpy.zeros(self.gm_.numberOfVariables,dtype=label_type)
          self.value_ = self.value_=self.gm_.evaluate(self.arg_)

        # start inference
        if visitor is not None:
          visitor.begin(self)

        if(self.initType=='localOpt'):
          print "move local opt"
          self.arg_ = self.gm_.moveLocalOpt('minimizer')
          print "done"
          visitor.visit(self)

        # end inference
        if visitor is not None:
          visitor.end(self)

    def name(self):
      return "CheapInitialization"

    def bound(self):
      return -1.0*float('inf')

    def arg(self):
        return self.arg_

    def value(self):
        return self.value_

inference.__dict__['CheapInitialization']=__CheapInitialization__
inference.__dict__['RandomFusion']=__RandomFusion__

if __name__ == "__main__":
    pass<commit_msg>fixed bug in saveGm function on python side<commit_after>#from opengmcore import _opengmcore.adder as adder
from opengmcore   import *
from __version__                    import version
from functionhelper                 import *
from _inf_param                     import _MetaInfParam , InfParam
from _visu                          import visualizeGm
from _misc                          import defaultAccumulator

from __version__ import version
import time

from _inference_interface_generator import _inject_interface , InferenceBase

import inference
import hdf5
import benchmark

# initialize solver/ inference dictionaries
_solverDicts=[
   (inference.adder.minimizer.solver.__dict__ ,     'adder',       'minimizer' ),
   (inference.adder.maximizer.solver.__dict__,      'adder',       'maximizer' ),
   (inference.multiplier.integrator.solver.__dict__,'adder',       'integrator'),
   (inference.multiplier.minimizer.solver.__dict__, 'multiplier',  'minimizer' ),
   (inference.multiplier.maximizer.solver.__dict__, 'multiplier',  'maximizer' ),
   (inference.multiplier.integrator.solver.__dict__,'multiplier',  'integrator')
]
for infClass,infName in _inject_interface(_solverDicts):
  inference.__dict__[infName]=infClass

class Timer(object):
    def __init__(self, name=None):
        self.name = name

    def __enter__(self):
        if self.name:
            print '[%s]' % self.name
        self.tstart = time.time()

    def __exit__(self, type, value, traceback):
        #if self.name:
        #    print '[%s]' % self.name,
        print '   Elapsed: %s' % (time.time() - self.tstart)

def saveGm(gm, f, d='gm'):
  """ save a graphical model to a hdf5 file:
  Args:
    gm : graphical model to save
    f  : filepath
    g  : dataset (defaut : 'gm')
  """
  hdf5.saveGraphicalModel(gm, f, d)

def loadGm(f, d='gm', operator='adder'):
  """ save a graphical model to a hdf5 file:
  Args:
    f  : filepath
    g  : dataset (defaut : 'gm')
    operator : operator of the graphical model ('adder' / 'multiplier')
  """
  if(operator=='adder'):
    gm=adder.GraphicalModel()
  elif(operator=='multiplier'):
    gm=multiplier.GraphicalModel()
  else:
    raise RuntimeError("unknown operator: "+ operator)
  hdf5.loadGraphicalModel(gm,f,d)
  return gm

class TestModels(object):
  @staticmethod
  def chain3(nVar,nLabels):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-2):
      f=numpy.random.rand(nLabels,nLabels,nLabels)
      model.addFactor(model.addFunction(f),[x0,x0+1,x0+2])
    return model

  @staticmethod
  def chain4(nVar,nLabels):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-3):
      f=numpy.random.rand(nLabels,nLabels,nLabels,nLabels)
      model.addFactor(model.addFunction(f),[x0,x0+1,x0+2,x0+3])
    return model

  @staticmethod
  def chainN(nVar,nLabels,order,nSpecialUnaries=0,beta=1.0):
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)

    for sn in range(nSpecialUnaries):
      r=int(numpy.random.rand(1)*nVar-1)
      rl=int(numpy.random.rand(1)*nLabels-1)

      unaries[r,rl]=0.0

    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    numpy.random.seed(42)
    for x0 in range(nVar-(order-1)):
      f=numpy.random.rand( *([nLabels]*order))
      f*=beta
      vis=numpy.arange(order)
      vis+=x0

      model.addFactor(model.addFunction(f),vis)
    return model

  @staticmethod
  def secondOrderGrid(dx,dy,nLabels):
    nVar=dx*dy
    model=adder.GraphicalModel([nLabels]*nVar)
    unaries = numpy.random.rand(nVar,nLabels)
    model.addFactors(model.addFunctions(unaries),numpy.arange(nVar))

    vis2Order=secondOrderGridVis(dx,dy,True)

    nF2=len(vis2Order)#.shape[0]
    f2s=numpy.random.rand(nF2,nLabels)

    model.addFactors(model.addFunctions(f2s),vis2Order)

    return model

class GenericTimingVisitor(object):
    def __init__(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        self.visitNth=visitNth
        self.reserve=reserve
        self.verbose=verbose
        self.multiline=multiline

        self.values_     = None
        self.runtimes_   = None
        self.bounds_     = None
        self.iterations_ = None

        self.t0          = None
        self.t1          = None
        self.iterNr      = 0
    def getValues(self):
        return numpy.require(self.values_,dtype=value_type)
    def getTimes(self):
        return numpy.require(self.runtimes_,dtype=value_type)
    def getBounds(self):
        return numpy.require(self.bounds_,dtype=value_type)
    def getIterations(self):
        return numpy.require(self.iterations_,dtype=value_type)

    def begin(self,inf):
      v = inf.value()
      b = inf.bound()

      self.values_    =[v]
      self.bounds_    =[b]
      self.runtimes_  =[0.0]
      self.iterations_=[self.iterNr]
      if self.verbose :
        print 'Begin :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)

      # start the timing
      self.t0         =time.time()
      self.t1         =time.time()
    def visit(self,inf):

      if(self.iterNr==0 or self.iterNr%self.visitNth==0):

        # "stop the timing"
        self.t1=time.time()

        # get the runtime of the run
        rt=self.t1-self.t0
        v = inf.value()
        b = inf.bound()
        if self.verbose :
          print 'Step  :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)

        # store results
        self.values_.append(v)
        self.bounds_.append(b)
        self.runtimes_.append(rt)
        self.iterations_.append(self.iterNr)

        # increment iteration number
        self.iterNr+=1

        # restart the timing
        self.t0=time.time()

      else:
        # increment iteration number
        self.iterNr+=1

    def end(self,inf):
        # "stop the timing"
        self.t1=time.time()

        # get the runtime of the run
        rt=self.t1-self.t0
        v = inf.value()
        b = inf.bound()
        if self.verbose :
          print 'End :        %d  Value : %f  Bound : %f '%(self.iterNr,v,b)
        # store results
        self.values_.append(v)
        self.bounds_.append(b)
        self.runtimes_.append(rt)
        self.iterations_.append(self.iterNr)

class __RandomFusion__(object):
    def __init__(self,gm,accumulator=None,parameter=InfParam()):

        if accumulator is None:
            self.accumulator=defaultAccumulator(gm=gm)
        else:
            self.accumulator=accumulator
        kwargs=parameter.kwargs
        self.gm_=gm

        self.steps = kwargs.get('steps', 100)
        self.fusionSolver = kwargs.get('fuisionSolver', 'lf2')

        self.arg_  = None
        self.value_ = None

        self.fusionMover=inference.adder.minimizer.FusionMover(self.gm_)

        self.nLabels = self.gm_.numberOfLabels(0)
        self.nVar    = self.gm_.numberOfVariables

    def timingVisitor(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        return GenericTimingVisitor(visitNth,reserve,verbose,multiline)

    def setStartingPoint(self,arg):
      self.arg_=arg
      self.value_=gm.evaluate(self.arg_)

    def infer(self,visitor=None):
        if(self.arg_ is None):
          self.arg_ = numpy.zeros(self.gm_.numberOfVariables,dtype=label_type)
          self.value_ = self.value_=self.gm_.evaluate(self.arg_)

        # start inference
        if visitor is not None:
          visitor.begin(self)

        # start fusion moves
        for x in range(self.steps):
          randState=numpy.random.randint(low=0, high=self.nLabels, size=self.nVar).astype(label_type)
          r = self.fusionMover.fuse(self.arg_,randState,self.fusionSolver)
          self.arg_=r[0]
          self.value_=r[1]
          visitor.visit(self)

        # end inference
        if visitor is not None:
          visitor.end(self)

    def name(self):
      return "RandomFusion"

    def bound(self):
      return -1.0*float('inf')

    def arg(self):
        return self.arg_

    def value(self):
        return self.value_

class __CheapInitialization__(object):
    def __init__(self,gm,accumulator=None,parameter=InfParam()):

        if accumulator is None:
            self.accumulator=defaultAccumulator(gm=gm)
        else:
            self.accumulator=accumulator
        kwargs=parameter.kwargs
        self.gm_=gm

        self.arg_  = None
        self.value_ = None
        self.initType = kwargs.get('initType', 'localOpt')

    def timingVisitor(self,visitNth=1,reserve=0,verbose=True,multiline=True):
        return GenericTimingVisitor(visitNth,reserve,verbose,multiline)

    def setStartingPoint(self,arg):
      self.arg_=arg
      self.value_=gm.evaluate(self.arg_)

    def infer(self,visitor=None):
        if(self.arg_ is None):
          self.arg_ = numpy.zeros(self.gm_.numberOfVariables,dtype=label_type)
          self.value_ = self.value_=self.gm_.evaluate(self.arg_)

        # start inference
        if visitor is not None:
          visitor.begin(self)

        if(self.initType=='localOpt'):
          print "move local opt"
          self.arg_ = self.gm_.moveLocalOpt('minimizer')
          print "done"
          visitor.visit(self)

        # end inference
        if visitor is not None:
          visitor.end(self)

    def name(self):
      return "CheapInitialization"

    def bound(self):
      return -1.0*float('inf')

    def arg(self):
        return self.arg_

    def value(self):
        return self.value_

inference.__dict__['CheapInitialization']=__CheapInitialization__
inference.__dict__['RandomFusion']=__RandomFusion__

if __name__ == "__main__":
    pass<|endoftext|>


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

from __future__ import print_function

import sqlite3
import psycopg2

def ir_sequence_export_sqlite(client, args, db_path, table_name, conn_string):

    conn = sqlite3.connect(db_path)
    conn.text_factory = str

    cursor = conn.cursor()
    try:
        cursor.execute('''DROP TABLE ''' + table_name + ''';''')
    except Exception as e:
        print('------->', e)
    cursor.execute(
        '''
        CREATE TABLE ''' + table_name + ''' (
            id INTEGER NOT NULL PRIMARY KEY,
            name,
            implementation,
            code,
            active,
            company_id,
            prefix,
            suffix,
            use_date_range,
            padding,
            number_increment,
            number_next,
            postgres_last_value,
            new_id INTEGER
            );
        '''
    )

    pg_conn = psycopg2.connect(conn_string)
    pg_cursor = pg_conn.cursor()

    client.context = {'active_test': False}
    ir_sequence_model = client.model('ir.sequence')
    ir_sequence_browse = ir_sequence_model.browse(args)

    ir_sequence_count = 0
    for ir_sequence_reg in ir_sequence_browse:
        ir_sequence_count += 1

        query = "SELECT last_value, increment_by, is_called FROM ir_sequence_%03d" % ir_sequence_reg.id
        pg_cursor.execute(query)
        row = pg_cursor.fetchone()
        postgres_last_value = row[0]

        print(
            ir_sequence_count, ir_sequence_reg.id, ir_sequence_reg.name.encode("utf-8"), ir_sequence_reg.code,
            postgres_last_value,
        )

        prefix = None
        if ir_sequence_reg.prefix:
            prefix = ir_sequence_reg.prefix

        suffix = None
        if ir_sequence_reg.suffix:
            suffix = ir_sequence_reg.suffix

        suffix = None
        if ir_sequence_reg.suffix:
            suffix = ir_sequence_reg.suffix

        cursor.execute('''
            INSERT INTO ''' + table_name + '''(
                id,
                name,
                implementation,
                code,
                active,
                company_id,
                prefix,
                suffix,
                use_date_range,
                padding,
                number_increment,
                number_next,
                postgres_last_value
                )
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)
            ''', (ir_sequence_reg.id,
                  ir_sequence_reg.name,
                  ir_sequence_reg.implementation,
                  ir_sequence_reg.code,
                  ir_sequence_reg.active,
                  ir_sequence_reg.company_id.id,
                  prefix,
                  suffix,
                  ir_sequence_reg.use_date_range,
                  ir_sequence_reg.padding,
                  ir_sequence_reg.number_increment,
                  ir_sequence_reg.number_next,
                  postgres_last_value,
                  )
        )

    conn.commit()
    conn.close()

    print()
    print('--> ir_sequence_count: ', ir_sequence_count)

def ir_sequence_import_sqlite(client, args, db_path, table_name, conn_string):

    # ir_sequence_model = client.model('ir.sequence')

    conn = sqlite3.connect(db_path)
    # conn.text_factory = str
    conn.row_factory = sqlite3.Row

    cursor = conn.cursor()

    data = cursor.execute('''
        SELECT
            id,
            name,
            implementation,
            code,
            active,
            company_id,
            prefix,
            suffix,
            use_date_range,
            padding,
            number_increment,
            number_next,
            postgres_last_value,
            new_id
        FROM ''' + table_name + ''';
    ''')

    print(data)
    print([field[0] for field in cursor.description])

    pg_conn = psycopg2.connect(conn_string)
    pg_cursor = pg_conn.cursor()

    ir_sequence_count = 0
    for row in cursor:
        ir_sequence_count += 1

        print(
            ir_sequence_count, row['id'], row['name'], row['implementation'], row['code'],
            row['prefix'], row['padding'], row['number_increment'], row['number_next'], row['postgres_last_value']
        )

        query = "SELECT last_value, increment_by, is_called FROM ir_sequence_%03d" % row['id']
        pg_cursor.execute(query)
        row2 = pg_cursor.fetchone()
        postgres_last_value = row2[0]

        my_table = row['code'][:-5].replace('.', '_')
        query2 = "SELECT MAX(id) FROM %s;" % my_table
        pg_cursor.execute(query2)
        row3 = pg_cursor.fetchone()
        max_id = row3[0]

        print('>>>>>', my_table, row['postgres_last_value'], postgres_last_value, max_id)

        if postgres_last_value < max_id:
            if row['postgres_last_value'] > max_id:
                query3 = "SELECT setval('ir_sequence_%03d', %d);" % (row['id'], row['postgres_last_value'])
                print('>>>>>>>>>>>>>>>', query3)
                pg_cursor.execute(query3)
                row4 = pg_cursor.fetchone()
                print('>>>>>>>>>>>>>>>', row4[0])
            else:
                query3 = "SELECT setval('ir_sequence_%03d', %d);" % (row['id'], max_id)
                print('>>>>>>>>>>>>>>>', query3)
                pg_cursor.execute(query3)
                row4 = pg_cursor.fetchone()
                print('>>>>>>>>>>>>>>>', row4[0])

    conn.commit()
    conn.close()

    print()
    print('--> ir_sequence_count: ', ir_sequence_count)
<commit_msg>[IMP] ir_sequence.py<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
###############################################################################
#
# Copyright (C) 2013-Today  Carlos Eduardo Vercelino - CLVsol
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################

from __future__ import print_function

import sqlite3
import psycopg2

def ir_sequence_export_sqlite(client, args, db_path, table_name, conn_string):

    conn = sqlite3.connect(db_path)
    conn.text_factory = str

    cursor = conn.cursor()
    try:
        cursor.execute('''DROP TABLE ''' + table_name + ''';''')
    except Exception as e:
        print('------->', e)
    cursor.execute(
        '''
        CREATE TABLE ''' + table_name + ''' (
            id INTEGER NOT NULL PRIMARY KEY,
            name,
            implementation,
            code,
            active,
            company_id,
            prefix,
            suffix,
            use_date_range,
            padding,
            number_increment,
            number_next,
            postgres_last_value,
            new_id INTEGER
            );
        '''
    )

    pg_conn = psycopg2.connect(conn_string)
    pg_cursor = pg_conn.cursor()

    client.context = {'active_test': False}
    ir_sequence_model = client.model('ir.sequence')
    ir_sequence_browse = ir_sequence_model.browse(args)

    ir_sequence_count = 0
    for ir_sequence_reg in ir_sequence_browse:
        ir_sequence_count += 1

        query = "SELECT last_value, increment_by, is_called FROM ir_sequence_%03d" % ir_sequence_reg.id
        pg_cursor.execute(query)
        row = pg_cursor.fetchone()
        postgres_last_value = row[0]

        print(
            ir_sequence_count, ir_sequence_reg.id, ir_sequence_reg.name.encode("utf-8"), ir_sequence_reg.code,
            postgres_last_value,
        )

        prefix = None
        if ir_sequence_reg.prefix:
            prefix = ir_sequence_reg.prefix

        suffix = None
        if ir_sequence_reg.suffix:
            suffix = ir_sequence_reg.suffix

        suffix = None
        if ir_sequence_reg.suffix:
            suffix = ir_sequence_reg.suffix

        cursor.execute('''
            INSERT INTO ''' + table_name + '''(
                id,
                name,
                implementation,
                code,
                active,
                company_id,
                prefix,
                suffix,
                use_date_range,
                padding,
                number_increment,
                number_next,
                postgres_last_value
                )
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)
            ''', (ir_sequence_reg.id,
                  ir_sequence_reg.name,
                  ir_sequence_reg.implementation,
                  ir_sequence_reg.code,
                  ir_sequence_reg.active,
                  ir_sequence_reg.company_id.id,
                  prefix,
                  suffix,
                  ir_sequence_reg.use_date_range,
                  ir_sequence_reg.padding,
                  ir_sequence_reg.number_increment,
                  ir_sequence_reg.number_next,
                  postgres_last_value,
                  )
        )

    conn.commit()
    conn.close()

    print()
    print('--> ir_sequence_count: ', ir_sequence_count)

def ir_sequence_import_sqlite(client, args, db_path, table_name, conn_string):

    # ir_sequence_model = client.model('ir.sequence')

    conn = sqlite3.connect(db_path)
    # conn.text_factory = str
    conn.row_factory = sqlite3.Row

    cursor = conn.cursor()

    data = cursor.execute('''
        SELECT
            id,
            name,
            implementation,
            code,
            active,
            company_id,
            prefix,
            suffix,
            use_date_range,
            padding,
            number_increment,
            number_next,
            postgres_last_value,
            new_id
        FROM ''' + table_name + ''';
    ''')

    print(data)
    print([field[0] for field in cursor.description])

    pg_conn = psycopg2.connect(conn_string)
    pg_cursor = pg_conn.cursor()

    ir_sequence_count = 0
    for row in cursor:
        ir_sequence_count += 1

        print(
            ir_sequence_count, row['id'], row['name'], row['implementation'], row['code'],
            row['prefix'], row['padding'], row['number_increment'], row['number_next'], row['postgres_last_value']
        )

        query = "SELECT last_value, increment_by, is_called FROM ir_sequence_%03d" % row['id']
        pg_cursor.execute(query)
        row2 = pg_cursor.fetchone()
        postgres_last_value = row2[0]

        my_table = row['code'][:-5].replace('.', '_')
        query2 = "SELECT MAX(id) FROM %s;" % my_table
        pg_cursor.execute(query2)
        row3 = pg_cursor.fetchone()
        max_id = row3[0]

        print('>>>>>', my_table, row['postgres_last_value'], postgres_last_value, max_id)

        if postgres_last_value < max_id:
            if row['postgres_last_value'] > max_id:
                query3 = "SELECT setval('ir_sequence_%03d', %d);" % (row['id'], row['postgres_last_value'])
                print('>>>>>>>>>>>>>>>', query3)
                pg_cursor.execute(query3)
                row4 = pg_cursor.fetchone()
                print('>>>>>>>>>>>>>>>', row4[0])
            else:
                query3 = "SELECT setval('ir_sequence_%03d', %d);" % (row['id'], max_id)
                print('>>>>>>>>>>>>>>>', query3)
                pg_cursor.execute(query3)
                row4 = pg_cursor.fetchone()
                print('>>>>>>>>>>>>>>>', row4[0])

    conn.commit()
    conn.close()

    print()
    print('--> ir_sequence_count: ', ir_sequence_count)

def ir_sequence_import_sqlite_named(client, args, db_path, table_name, conn_string, input_code, output_code):

    # ir_sequence_model = client.model('ir.sequence')

    conn = sqlite3.connect(db_path)
    # conn.text_factory = str
    conn.row_factory = sqlite3.Row

    cursor = conn.cursor()

    data = cursor.execute('''
        SELECT
            id,
            name,
            implementation,
            code,
            active,
            company_id,
            prefix,
            suffix,
            use_date_range,
            padding,
            number_increment,
            number_next,
            postgres_last_value,
            new_id
        FROM ''' + table_name + ''';
    ''')

    print(data)
    print([field[0] for field in cursor.description])

    pg_conn = psycopg2.connect(conn_string)
    pg_cursor = pg_conn.cursor()

    ir_sequence_model = client.model('ir.sequence')

    ir_sequence_count = 0
    for row in cursor:
        ir_sequence_count += 1

        if input_code == row['code']:

            print(
                ir_sequence_count, row['id'], row['name'], row['implementation'], row['code'],
                row['prefix'], row['padding'], row['number_increment'], row['number_next'], row['postgres_last_value']
            )

            ir_sequence_browse = ir_sequence_model.browse(args)

            ir_sequence_count = 0
            for ir_sequence_reg in ir_sequence_browse:
                ir_sequence_count += 1

                if ir_sequence_reg.code == output_code:

                    query = "SELECT last_value, increment_by, is_called FROM ir_sequence_%03d" % ir_sequence_reg.id
                    pg_cursor.execute(query)
                    row5 = pg_cursor.fetchone()
                    postgres_last_value = row5[0]

                    print(
                        ir_sequence_count, ir_sequence_reg.id, ir_sequence_reg.name.encode("utf-8"),
                        ir_sequence_reg.code,
                        postgres_last_value,
                    )

                    my_table = ir_sequence_reg.code[:-5].replace('.', '_')
                    query2 = "SELECT MAX(id) FROM %s;" % my_table
                    pg_cursor.execute(query2)
                    row3 = pg_cursor.fetchone()
                    max_id = row3[0]

                    print('>>>>>', my_table, row['postgres_last_value'], postgres_last_value, max_id,
                          ir_sequence_reg.number_next_actual)

                    if postgres_last_value < max_id:
                        if row['postgres_last_value'] > max_id:
                            query3 = "SELECT setval('ir_sequence_%03d', %d);" % \
                                (ir_sequence_reg.id, row['postgres_last_value'])
                            print('>>>>>>>>>>>>>>>', query3)
                            pg_cursor.execute(query3)
                            row4 = pg_cursor.fetchone()
                            print('>>>>>>>>>>>>>>>', row4[0])
                        else:
                            query3 = "SELECT setval('ir_sequence_%03d', %d);" % (ir_sequence_reg.id, max_id)
                            print('>>>>>>>>>>>>>>>', query3)
                            pg_cursor.execute(query3)
                            row4 = pg_cursor.fetchone()
                            print('>>>>>>>>>>>>>>>', row4[0])
                    else:
                        query3 = "SELECT setval('ir_sequence_%03d', %d);" % (ir_sequence_reg.id, postgres_last_value)
                        print('>>>>>>>>>>>>>>>', query3)
                        pg_cursor.execute(query3)
                        row4 = pg_cursor.fetchone()
                        print('>>>>>>>>>>>>>>>', row4[0])

    conn.commit()
    conn.close()

    print()
    print('--> ir_sequence_count: ', ir_sequence_count)
<|endoftext|>