<commit_before>"""
A simulated manipulator based upon Baxter robot
to write given letter trajectory
"""

import os

import sys
import copy

from collections import defaultdict

import numpy as np
import matplotlib.pyplot as plt

import baxter_pykdl

import pylqr.pylqr_trajctrl as plqrtc

import utils

class BaxterWriter():
    def __init__(self):
        self.prepare_baxter_robot_manipulator(manip_idx=1)

        #the center of block to write the letter... in the world reference frame
        self.block_center = np.array([0.844, -0.357, 0.257])
        self.scale = 0.09 / (1.5 + 1.5) #boundary of block / dim of trajectories
        #this is for the right arm
        self.seed_pos = [0.00230, -0.77274, 0.95529, 1.53091, -0.91924, 0.55914, -0.09664]
        self.seed_pose = self.robot_dynamics.forward_position_kinematics(self.seed_pos)
        return

    def prepare_baxter_robot_manipulator(self, manip_idx=1):
        path_prefix = os.path.dirname(os.path.abspath(__file__))
        #urdf
        self.baxter_urdf_path = os.path.join(path_prefix, 'urdf/baxter.urdf')

        #default is the right one: 1
        if manip_idx == 0:
            kin_name = 'left'
        else:
            kin_name = 'right'
        # use revised baxter_pykdl to create inverse kinemtics model
        self.robot_dynamics = baxter_pykdl.baxter_pykdl.baxter_dynamics(kin_name, self.baxter_urdf_path)
        #print structure
        self.robot_dynamics.print_robot_description()

        return

    def generate_spatial_trajectory(self, char_traj):
        """
        Member function to generate spatial trajectory from 2D char trajectory through translation/rotation/scaling
        """
        #char_traj is supposed to be a 2D array
        spatial_traj = np.zeros((len(char_traj), 3))
        spatial_traj[:, 0:2] = char_traj * self.scale
        spatial_traj = spatial_traj + self.block_center
        return spatial_traj

    def derive_ik_trajectory(self, spatial_traj):
        # get orientation from the seed pose
        q_seed = self.seed_pos
        q_array = []
        for idx, pnt in enumerate(spatial_traj):
            #solve IK
            q = self.robot_dynamics.inverse_kinematics(pnt, orientation=self.seed_pose[3:7], seed=q_seed)
            if q is None:
                print 'Failed to solve IK at step {0} for desired position {1}.'.format(idx, pnt)
            else:
                q_seed = q
                q_array.append(q)
        return np.array(q_array)

    def derive_cartesian_trajectory(self, q_array):
        #derive cartesian trajectory from q_array
        spatial_traj = []
        for idx, q in enumerate(q_array):
            cart_pose = self.robot_dynamics.forward_position_kinematics(q)
            spatial_traj.append(cart_pose)
        return spatial_traj

    def derive_ilqr_trajectory(self, spatial_traj):
        dt = .01
        #derive joint trajectory from ilqr
        lqr_traj_ctrl = plqrtc.PyLQR_TrajCtrl(R=.01, dt=dt)

        #cost function
        def traj_cost(x, u, t, aux):
            #x is the given joint position, evaluate forward kinematics
            cart_pose = self.robot_dynamics.forward_position_kinematics(x[0:7])
            track_err = np.linalg.norm((cart_pose[0:3] - spatial_traj[t])*np.array([10., 10., 50]))**2
            #control effort from inverse dynamics
            jnt_mass = self.robot_dynamics.inertia(x[:7])
            # jnt_coriolis = self.robot_dynamics.coriolis(x[:7], x[7:])
            # jnt_gravity = self.robot_dynamics.gravity(x[:7])
            # tau = jnt_mass.dot(u) + jnt_coriolis + jnt_gravity
            # tau = jnt_coriolis + jnt_gravity
            tau = jnt_mass.dot(u)
            #control effort from control input
            # tau = u
            ctrl_effort = np.linalg.norm(tau) ** 2
            return track_err + lqr_traj_ctrl.R_ * ctrl_effort

        lqr_traj_ctrl.build_ilqr_general_solver(cost_func=traj_cost, n_dims=7, T=100)

        #prepare initial guess of trajectory, from the IK solution
        q_ik = self.derive_ik_trajectory(spatial_traj)
        x0 = q_ik[0] #init velocity is zero, this is handled in the ilqr_traj_ctrl

        q_dot_ik = np.diff(q_ik, axis=0) / dt
        q_dot_ik = np.vstack([np.zeros(7), q_dot_ik])   #initial velocity is zero
        q_ddot_ik = np.diff(q_dot_ik, axis=0) / dt
        u_array_init = q_ddot_ik
        #test this init
        #note there will be some difference as the trajectory ilqr force the initial velocity as zero
        #while the finite difference is not the case
        # x_array = lqr_traj_ctrl.ilqr_.forward_propagation(np.concatenate([x0, np.zeros(7)]), u_array_init)
        # print q_ik - np.array(x_array)[:, 0:7]

        syn_traj = lqr_traj_ctrl.synthesize_trajectory(x0, u_array_init, n_itrs=25)
        return syn_traj[:, 0:7]

def build_ik_joint_traj_for_chars(data):
    baxter_writer = BaxterWriter()
    res_data = defaultdict(list)

    for c in data.keys():
        print 'Processing character {0}...'.format(c)
        for d in data[c]:
            tmp_char_traj = np.reshape(d[:-1], (2, -1)).T
            tmp_spatial_traj = baxter_writer.generate_spatial_trajectory(tmp_char_traj)

            tmp_q_array = baxter_writer.derive_ik_trajectory(tmp_spatial_traj)

            #note the data is transposed and flattened as the cartesian trajectories
            res_data[c].append(np.array(tmp_q_array).T.flatten())

    return res_data

def check_joint_ik_data(cart_data, jnt_data, n_chars=5, n_samples=1):
    baxter_writer = BaxterWriter()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hold(True)
    plt.ion()

    check_chars = [np.random.choice(cart_data.keys()) for i in range(n_chars)]
    for c in check_chars:
        check_indices = [np.random.choice(range(len(cart_data[c]))) for i in range(n_samples)]
        #see how it's going for all the samples
        for idx in check_indices:
            tmp_char_traj = np.reshape(cart_data[c][idx][:-1], (2, -1)).T
            tmp_spatial_traj = baxter_writer.generate_spatial_trajectory(tmp_char_traj)

            ax.plot(tmp_spatial_traj[:, 0], tmp_spatial_traj[:, 1], 'k*', linewidth=3.5)

            #reconstruction from joint trajectories...
            tmp_jnt_traj = np.reshape(jnt_data[c][idx], (7, -1)).T
            tmp_cart_array = baxter_writer.derive_cartesian_trajectory(tmp_jnt_traj)
            recons_char_traj = np.array([cart_pose[0:2] for cart_pose in tmp_cart_array])
            z_array = [cart_pose[3] for cart_pose in tmp_cart_array]
            print 'Z - mean and std:', np.mean(z_array), np.std(z_array)

            #show the reconstructed trajectory
            ax.plot(recons_char_traj[:, 0], recons_char_traj[:, 1], 'r', linewidth=3.5)

            plt.draw()
    return

def main():
    baxter_writer = BaxterWriter()

    #prepare a circular path
    n_pnts = 100
    t = np.linspace(0, 2*np.pi, n_pnts)
    a = 1.35
    b = 1.0
    char_traj = np.array([a*np.cos(t), b*np.sin(t)]).T

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hold(True)
    #derive joint trajectory
    spatial_traj = baxter_writer.generate_spatial_trajectory(char_traj)

    ax.plot(spatial_traj[:, 0], spatial_traj[:, 1], 'k*', linewidth=3.5)
    print 'Solving a series of IK problems...'
    q_array = baxter_writer.derive_ik_trajectory(spatial_traj)
    print 'Finished solving IK problems.'
    #restore to cartesian motion
    print 'Reconstructing Cartesian trajectory...'
    cart_array = baxter_writer.derive_cartesian_trajectory(q_array)
    print 'Finished reconstructing Cartesian trajectory.'
    #extract 2D trajectory
    recons_char_traj = np.array([cart_pose[0:2] for cart_pose in cart_array])

    z_array = [cart_pose[2] for cart_pose in cart_array]
    print 'IK - Z - mean and std:', np.mean(z_array), np.std(z_array)

    #show the reconstructed trajectory
    ax.plot(recons_char_traj[:, 0], recons_char_traj[:, 1], 'r', linewidth=3.5)

    raw_input('ENTER to continue derive iLQR optimal control')

    print 'Solving iLQR optimal trajectories'
    q_array_ilqr = baxter_writer.derive_ilqr_trajectory(spatial_traj)
    print 'Finished solving iLQR optimal control'
    #restore to cartesian motion
    print 'Reconstructing Cartesian trajectory...'
    cart_array_ilqr = baxter_writer.derive_cartesian_trajectory(q_array_ilqr)
    print 'Finished reconstructing Cartesian trajectory.'
    #extract 2D trajectory
    recons_char_traj_ilqr = np.array([cart_pose[0:2] for cart_pose in cart_array_ilqr])
    z_array_ilqr = [cart_pose[2] for cart_pose in cart_array_ilqr]
    print 'iLQR - Z - mean and std:', np.mean(z_array_ilqr), np.std(z_array_ilqr)

    ax.plot(recons_char_traj_ilqr[:, 0], recons_char_traj_ilqr[:, 1], 'g', linewidth=3.5)

    plt.show()
    return

if __name__ == '__main__':
    main()<commit_msg>generate ilqr data; code seems to work fine.<commit_after>"""
A simulated manipulator based upon Baxter robot
to write given letter trajectory
"""

import os

import sys
import copy

from collections import defaultdict

import numpy as np
import matplotlib.pyplot as plt

import baxter_pykdl

import pylqr.pylqr_trajctrl as plqrtc

import utils

class BaxterWriter():
    def __init__(self):
        self.prepare_baxter_robot_manipulator(manip_idx=1)

        #the center of block to write the letter... in the world reference frame
        self.block_center = np.array([0.844, -0.357, 0.257])
        self.scale = 0.09 / (1.5 + 1.5) #boundary of block / dim of trajectories
        #this is for the right arm
        self.seed_pos = [0.00230, -0.77274, 0.95529, 1.53091, -0.91924, 0.55914, -0.09664]
        self.seed_pose = self.robot_dynamics.forward_position_kinematics(self.seed_pos)
        return

    def prepare_baxter_robot_manipulator(self, manip_idx=1):
        path_prefix = os.path.dirname(os.path.abspath(__file__))
        #urdf
        self.baxter_urdf_path = os.path.join(path_prefix, 'urdf/baxter.urdf')

        #default is the right one: 1
        if manip_idx == 0:
            kin_name = 'left'
        else:
            kin_name = 'right'
        # use revised baxter_pykdl to create inverse kinemtics model
        self.robot_dynamics = baxter_pykdl.baxter_pykdl.baxter_dynamics(kin_name, self.baxter_urdf_path)
        #print structure
        self.robot_dynamics.print_robot_description()

        return

    def generate_spatial_trajectory(self, char_traj):
        """
        Member function to generate spatial trajectory from 2D char trajectory through translation/rotation/scaling
        """
        #char_traj is supposed to be a 2D array
        spatial_traj = np.zeros((len(char_traj), 3))
        spatial_traj[:, 0:2] = char_traj * self.scale
        spatial_traj = spatial_traj + self.block_center
        return spatial_traj

    def derive_ik_trajectory(self, spatial_traj):
        # get orientation from the seed pose
        q_seed = self.seed_pos
        q_array = []
        for idx, pnt in enumerate(spatial_traj):
            #solve IK
            q = self.robot_dynamics.inverse_kinematics(pnt, orientation=self.seed_pose[3:7], seed=q_seed)
            if q is None:
                print 'Failed to solve IK at step {0} for desired position {1}.'.format(idx, pnt)
            else:
                q_seed = q
                q_array.append(q)
        return np.array(q_array)

    def derive_cartesian_trajectory(self, q_array):
        #derive cartesian trajectory from q_array
        spatial_traj = []
        for idx, q in enumerate(q_array):
            cart_pose = self.robot_dynamics.forward_position_kinematics(q)
            spatial_traj.append(cart_pose)
        return spatial_traj

    def derive_ilqr_trajectory(self, spatial_traj):
        dt = .01
        #derive joint trajectory from ilqr
        lqr_traj_ctrl = plqrtc.PyLQR_TrajCtrl(R=.01, dt=dt)

        #cost function
        def traj_cost(x, u, t, aux):
            #x is the given joint position, evaluate forward kinematics
            cart_pose = self.robot_dynamics.forward_position_kinematics(x[0:7])
            track_err = np.linalg.norm((cart_pose[0:3] - spatial_traj[t])*np.array([10., 10., 50]))**2
            #control effort from inverse dynamics
            jnt_mass = self.robot_dynamics.inertia(x[:7])
            # jnt_coriolis = self.robot_dynamics.coriolis(x[:7], x[7:])
            # jnt_gravity = self.robot_dynamics.gravity(x[:7])
            # tau = jnt_mass.dot(u) + jnt_coriolis + jnt_gravity
            # tau = jnt_coriolis + jnt_gravity
            tau = jnt_mass.dot(u)
            #control effort from control input
            # tau = u
            ctrl_effort = np.linalg.norm(tau) ** 2
            return track_err + lqr_traj_ctrl.R_ * ctrl_effort

        lqr_traj_ctrl.build_ilqr_general_solver(cost_func=traj_cost, n_dims=7, T=100)

        #prepare initial guess of trajectory, from the IK solution
        q_ik = self.derive_ik_trajectory(spatial_traj)
        x0 = q_ik[0] #init velocity is zero, this is handled in the ilqr_traj_ctrl

        q_dot_ik = np.diff(q_ik, axis=0) / dt
        q_dot_ik = np.vstack([np.zeros(7), q_dot_ik])   #initial velocity is zero
        q_ddot_ik = np.diff(q_dot_ik, axis=0) / dt
        u_array_init = q_ddot_ik
        #test this init
        #note there will be some difference as the trajectory ilqr force the initial velocity as zero
        #while the finite difference is not the case
        # x_array = lqr_traj_ctrl.ilqr_.forward_propagation(np.concatenate([x0, np.zeros(7)]), u_array_init)
        # print q_ik - np.array(x_array)[:, 0:7]

        syn_traj = lqr_traj_ctrl.synthesize_trajectory(x0, u_array_init, n_itrs=25)
        return syn_traj[:, 0:7]

def build_ik_joint_traj_for_chars(data):
    baxter_writer = BaxterWriter()
    res_data = defaultdict(list)

    for c in data.keys():
        print 'Processing character {0}...'.format(c)
        for d in data[c]:
            tmp_char_traj = np.reshape(d[:-1], (2, -1)).T
            tmp_spatial_traj = baxter_writer.generate_spatial_trajectory(tmp_char_traj)

            tmp_q_array = baxter_writer.derive_ik_trajectory(tmp_spatial_traj)

            #note the data is transposed and flattened as the cartesian trajectories
            res_data[c].append(np.array(tmp_q_array).T.flatten())

    return res_data

def build_ilqr_joint_traj_for_chars(data):
    baxter_writer = BaxterWriter()
    res_data = defaultdict(list)

    for c in data.keys():
        print 'Processing character {0}...'.format(c)
        for d in data[c]:
            tmp_char_traj = np.reshape(d[:-1], (2, -1)).T
            tmp_spatial_traj = baxter_writer.generate_spatial_trajectory(tmp_char_traj)

            tmp_q_array = baxter_writer.derive_ilqr_trajectory(tmp_spatial_traj)

            #note the data is transposed and flattened as the cartesian trajectories
            res_data[c].append(np.array(tmp_q_array).T.flatten())

    return res_data

def check_joint_data(cart_data, jnt_data, n_chars=5, n_samples=1):
    baxter_writer = BaxterWriter()
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hold(True)
    plt.ion()

    check_chars = [np.random.choice(cart_data.keys()) for i in range(n_chars)]
    for c in check_chars:
        check_indices = [np.random.choice(range(len(cart_data[c]))) for i in range(n_samples)]
        #see how it's going for all the samples
        for idx in check_indices:
            tmp_char_traj = np.reshape(cart_data[c][idx][:-1], (2, -1)).T
            tmp_spatial_traj = baxter_writer.generate_spatial_trajectory(tmp_char_traj)

            ax.plot(tmp_spatial_traj[:, 0], tmp_spatial_traj[:, 1], 'k*', linewidth=3.5)

            #reconstruction from joint trajectories...
            tmp_jnt_traj = np.reshape(jnt_data[c][idx], (7, -1)).T
            tmp_cart_array = baxter_writer.derive_cartesian_trajectory(tmp_jnt_traj)
            recons_char_traj = np.array([cart_pose[0:2] for cart_pose in tmp_cart_array])
            z_array = [cart_pose[3] for cart_pose in tmp_cart_array]
            print 'Z - mean and std:', np.mean(z_array), np.std(z_array)

            #show the reconstructed trajectory
            ax.plot(recons_char_traj[:, 0], recons_char_traj[:, 1], 'r', linewidth=3.5)

            plt.draw()
    return

def main():
    baxter_writer = BaxterWriter()

    #prepare a circular path
    n_pnts = 100
    t = np.linspace(0, 2*np.pi, n_pnts)
    a = 1.35
    b = 1.0
    char_traj = np.array([a*np.cos(t), b*np.sin(t)]).T

    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.hold(True)
    #derive joint trajectory
    spatial_traj = baxter_writer.generate_spatial_trajectory(char_traj)

    ax.plot(spatial_traj[:, 0], spatial_traj[:, 1], 'k*', linewidth=3.5)
    print 'Solving a series of IK problems...'
    q_array = baxter_writer.derive_ik_trajectory(spatial_traj)
    print 'Finished solving IK problems.'
    #restore to cartesian motion
    print 'Reconstructing Cartesian trajectory...'
    cart_array = baxter_writer.derive_cartesian_trajectory(q_array)
    print 'Finished reconstructing Cartesian trajectory.'
    #extract 2D trajectory
    recons_char_traj = np.array([cart_pose[0:2] for cart_pose in cart_array])

    z_array = [cart_pose[2] for cart_pose in cart_array]
    print 'IK - Z - mean and std:', np.mean(z_array), np.std(z_array)

    #show the reconstructed trajectory
    ax.plot(recons_char_traj[:, 0], recons_char_traj[:, 1], 'r', linewidth=3.5)

    raw_input('ENTER to continue derive iLQR optimal control')

    print 'Solving iLQR optimal trajectories'
    q_array_ilqr = baxter_writer.derive_ilqr_trajectory(spatial_traj)
    print 'Finished solving iLQR optimal control'
    #restore to cartesian motion
    print 'Reconstructing Cartesian trajectory...'
    cart_array_ilqr = baxter_writer.derive_cartesian_trajectory(q_array_ilqr)
    print 'Finished reconstructing Cartesian trajectory.'
    #extract 2D trajectory
    recons_char_traj_ilqr = np.array([cart_pose[0:2] for cart_pose in cart_array_ilqr])
    z_array_ilqr = [cart_pose[2] for cart_pose in cart_array_ilqr]
    print 'iLQR - Z - mean and std:', np.mean(z_array_ilqr), np.std(z_array_ilqr)

    ax.plot(recons_char_traj_ilqr[:, 0], recons_char_traj_ilqr[:, 1], 'g', linewidth=3.5)

    plt.show()
    return

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


<commit_before># -*- coding: utf-8 -*-
"""
Created on Tue Apr 19 16:31:18 2016

@author: imorgan.admin
"""

from bayespy.jni import bayesServerInference

import bayespy.network
import pandas as pd
from bayespy.jni import bayesServer
from bayespy.jni import bayesServerParams
from bayespy.jni import jp
import numpy as np
import logging
import pathos.multiprocessing as mp
from typing import List

class QueryOutput:
    def __init__(self, continuous, discrete):
        self.continuous = continuous
        self.discrete = discrete

class InferenceEngine:

    _inference_factory = None

    def __init__(self, network):
        self._network = network

    @staticmethod
    def get_inference_factory():
        if InferenceEngine._inference_factory is None:
            InferenceEngine._inference_factory = bayesServerInference().RelevanceTreeInferenceFactory()
        return InferenceEngine._inference_factory

    def create(self, loglikelihood=False, conflict=False, retract=False):
        query_options = self.get_inference_factory().createQueryOptions()
        query_output = self.get_inference_factory().createQueryOutput()
        inference_engine = self.get_inference_factory().createInferenceEngine(self._network)

        query_options.setLogLikelihood(loglikelihood)
        query_options.setConflict(conflict)

        if retract:
            query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)

        return inference_engine, query_options, query_output

class Query:

    def __init__(self, jnetwork, inference):
        self._factory = bayesServerInference().RelevanceTreeInferenceFactory()
        self._queryOptions = self._factory.createQueryOptions()
        self._queryOutput = self._factory.createQueryOutput()
        self._queryDistributions = inference.getQueryDistributions()
        self._inference = inference
        self._jnetwork = jnetwork

    def execute(self, variables=None):
        """
        Query a number of variables (if none, then query all variables in the network)
        :param variables: a list of variables, or none
        :return: a QueryOutput object with separate continuous/ discrete dataframes
        """
        distributions = {}
        if variables is None:
            variables = self._jnetwork.getVariables()

        for v in variables:
            if bayespy.network.is_variable_discrete(v):
                table = bayesServer().Table(v)
            else:
                table = bayesServer().CLGaussian(v)

            self._queryDistributions.add(bayesServerInference().QueryDistribution(table))
            distributions.update({ v.getName() : table})

        self._inference.query(self._queryOptions, self._queryOutput)

        states= []
        d_variables = []
        values = []
        c_variables = []
        mean = []
        variance = []
        for v in variables:
            dist =  distributions[v.getName()]
            if bayespy.network.is_variable_discrete(v):
                for state in v.getStates():
                    states.append(state.getName())
                    values.append(float(dist.get([state])))
                    d_variables.append(v.getName())
            else:
                variance.append(float(dist.getVariance(v)))
                mean.append(float(dist.getMean(v)))
                c_variables.append(v.getName())

        return QueryOutput(pd.DataFrame({ 'variable': c_variables, 'mean': mean, 'variance': variance}), pd.DataFrame({ 'variable' : d_variables, 'state': states, 'value' : values}))

class Evidence:

    def __init__(self, jnetwork, inference):
        self._jnetwork = jnetwork
        self._inference = inference
        self._evidence = inference.getEvidence()
        self._evidence.clear()
        self._variables = jnetwork.getVariables()

    def apply(self, evidence):
        """
        Apply evidence to a network
        :param evidence: if a discrete variable, then a list of strings generated by network.Discrete.tostring(),
        if continuous, then a list of tuples with (VariableName, ContinuousValue)
        :return: Nothing
        """
        for value in evidence:
            if not isinstance(value, tuple):
                value = str(value)
                node, state = value.split(bayespy.network.STATE_DELIMITER)
                v = self._variables.get(node)
                if v is None:
                    raise ValueError("Node {} does not exist".format(node))
                if bayespy.network.is_variable_discrete(v):
                    st = v.getStates().get(state)
                    if st is None:
                        raise ValueError("State {} does not exist in variable {}".format(state, node))

                    self._evidence.setState(st)
                else:
                    raise ValueError("Can not find state on a continuous variable")
            else:
                v = self._variables.get(value[0])
                if not bayespy.network.is_variable_continuous(v):
                    raise ValueError("Variable is not continuous")

                self._evidence.set(v, jp.java.lang.Double(value[1]))

        return self._evidence

class QueryStatistics:

    def __init__(self, calc_loglikelihood=True, calc_conflict=False, loglikelihood_column='loglikelihood', conflict_column='conflict'):
        self._calc_loglikelihood = calc_loglikelihood
        self._calc_conflict = calc_conflict
        self._loglikelihood_column = loglikelihood_column
        self._conflict_column = conflict_column

    def setup(self, network, inference_engine, query_options):
        query_options.setLogLikelihood(self._calc_loglikelihood)
        query_options.setConflict(self._calc_conflict)

    def results(self, inference_engine, query_output):
        result = {}
        if self._calc_loglikelihood:
            result.update({ self._loglikelihood_column: query_output.getLogLikelihood().floatValue() })

        if self._calc_conflict:
            result.update({ self._conflict_column: query_output.getConflict().floatValue()})

        return result

# seems like a better name than QueryStatistics, so just having this here.
class QueryModelStatistics(QueryStatistics):
    def __init__(self, calc_loglikelihood=True, calc_conflict=False, loglikelihood_column='loglikelihood', conflict_column='conflict'):
        super().__init__(calc_loglikelihood, calc_conflict, loglikelihood_column, conflict_column)

class QueryMostLikelyState:

    def __init__(self, target_variable_name, output_dtype="object", suffix="_maxlikelihood"):
        self._target_variable_name = target_variable_name
        self._distribution = None
        self._output_dtype = output_dtype
        self._suffix = suffix

    def setup(self, network, inference_engine, query_options):
        distribution = None

        self._variable = bayespy.network.get_variable(network, self._target_variable_name)

        if bayespy.network.is_variable_discrete(self._variable):
            distribution = bayesServer().Table(self._variable)

        if distribution is None:
            raise ValueError("{} needs to be discrete in QueryMostLikelyState".format(self._target_variable_name))

        query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)
        qd = bayesServerInference().QueryDistribution(distribution)

        self._distribution = distribution
        inference_engine.getQueryDistributions().add(qd)

    def results(self, inference_engine, query_output):
        states = {}
        for state in self._variable.getStates():
            states.update({ state.getName() : self._distribution.get([state])})

        # get the most likely state
        max_state = max(states.keys(), key=(lambda key: states[key]))
        max_state_name = bayespy.data.DataFrame.cast2(self._output_dtype, max_state)

        return {self._target_variable_name + self._suffix: max_state_name}

class QueryLogLikelihood:
    def __init__(self, variable_names, column_name: str='_loglikelihood'):
        if isinstance(variable_names, str):
            variable_names = [variable_names]

        self._variable_names = variable_names
        self._distribution = None
        self._query_distribution = None

        self._column_name = column_name

    def setup(self, network, inference_engine, query_options):
        variables = [bayespy.network.get_variable(network, n) for n in self._variable_names]
        if len(variables) == 1:
            self._distribution = bayesServer().CLGaussian(variables[0])
        else:
            self._distribution = bayesServer().CLGaussian(variables)

        query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)
        qd = bayesServerInference().QueryDistribution(self._distribution)
        qd.setQueryLogLikelihood(True)
        self._query_distribution = qd
        inference_engine.getQueryDistributions().add(qd)

    def results(self, inference_engine, query_output):
        result = {}
        ll = self._query_distribution.getLogLikelihood()
        value = ll.floatValue() if ll is not None else np.nan
        result.update({":".join(self._variable_names) + self._column_name: value})
        return result

class QueryMeanVariance:
    def __init__(self, variable_name, retract_evidence=True, result_mean_suffix='_mean', result_variance_suffix='_variance'):
        self._variable_name = variable_name

        self._result_mean_suffix = result_mean_suffix
        self._result_variance_suffix = result_variance_suffix
        self._retract_evidence = retract_evidence

    def setup(self, network, inference_engine, query_options):
        self._variable = bayespy.network.get_variable(network, self._variable_name)

        if not bayespy.network.is_variable_continuous(self._variable):
            raise ValueError("{} needs to be continuous.".format(self._variable_name))

        self._query = bayesServer().CLGaussian(self._variable)

        if self._retract_evidence:
            query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)

        inference_engine.getQueryDistributions().add(bayesServerInference().QueryDistribution(self._query))

    def results(self, inference_engine, query_output):
        return {self._variable_name + self._result_mean_suffix: self._query.getMean(self._variable),
                self._variable_name + self._result_variance_suffix: self._query.getVariance(self._variable)}

def _batch_query(df: pd.DataFrame, connection_string: str, network: str, table_name: str,
                 variable_references: List[str],
                 queries, logger, i):

    bayespy.jni.attach_thread()
    network = bayespy.network.create_network_from_string(network)

    reader_options = bayesServer().data.ReaderOptions("ix")
    variable_refs = list(bayespy.network.create_variable_references(network, df,
                                                                    variable_references=variable_references))

    data_reader = bayesServer().data.DatabaseDataReaderCommand(
        connection_string,
        "select * from {} where ix in ({})".format(table_name, ",".join(str(i) for i in df.index.tolist()))).executeReader()

    reader = bayesServer().data.DefaultEvidenceReader(data_reader, jp.java.util.Arrays.asList(variable_refs),
                                                   reader_options)

    factory = InferenceEngine(network)
    (inference_engine, query_options, query_output) = factory.create()

    for query in queries:
        query.setup(network, inference_engine, query_options)

    results = []
    try:
        while reader.read(inference_engine.getEvidence(), bayesServer().data.DefaultReadOptions(True)):
            result = {}

            try:
                inference_engine.query(query_options, query_output)
            except BaseException as e:
                logger.error(e)
                # inference_engine.getEvidence().clear()
                # continue

            for query in queries:
                result = {**result, **query.results(inference_engine, query_output)}

            inference_engine.getEvidence().clear()
            result.update({'caseid': int(reader.getReadInfo().getCaseId().toString())})

            results.append(result)

            if i % 500 == 0:
                logger.info("Queried case {}".format(i))

            i += 1
    finally:
        reader.close()
        #bayespy.jni.detach()
        return results

class BatchQuery:

    def __init__(self, network, datastore, logger: logging.Logger):

        self._logger = logger
        self._datastore = datastore
        # serialise the network as a string.
        self._network = network.saveToString()

    def _calc_num_threads(self, df_size: int, query_size: int) -> int:
        num_queries = df_size * query_size

        if mp.cpu_count() == 2:
            max = 2
        else:
            max = mp.cpu_count() - 1

        calc = int(num_queries / 5000)
        if calc > max:
            r = max
        elif calc <= 1:
            if num_queries > 1000:
                r = 2
            else:
                r = 1
        else:
            r = calc

        if r <= 1:
            r = 1

        return r

    def query(self, queries=[QueryStatistics()], append_to_df=True, variable_references=[]):

        if not hasattr(queries, "__getitem__"):
            queries = [queries]

        nt = self._network
        logger = self._logger
        conn = self._datastore.get_connection()
        table = self._datastore.table
        processes = self._calc_num_threads(len(self._datastore.data), len(queries))

        self._logger.info("Using {} processes to query {} rows".format(processes, len(self._datastore.data)))

        if processes == 1:
            pdf = pd.DataFrame(_batch_query(self._datastore.data, conn, nt, table,
                         variable_references, queries,
                         logger, 0))
        else:
            with mp.Pool(processes=processes) as pool:
                pdf = pd.DataFrame()
                for result_set in pool.map(lambda df: _batch_query(df, conn, nt, table,
                                                                 variable_references, queries,
                                                                logger, 0), np.array_split(self._datastore.data, processes)):
                    pdf = pdf.append(pd.DataFrame(result_set))

        df = pdf.set_index('caseid')

        if append_to_df:
            return self._datastore.data.join(df)
        else:
            return df

class NetworkModel:

    def __init__(self, network, data_store, logger):
        self._jnetwork = network
        self._inference_factory = InferenceEngine(network)
        self._data_store = data_store
        self._logger = logger
        self._data = data_store.data

    def get_network(self):
        return self._jnetwork

    def save(self, path):
        from xml.dom import minidom
        nt = self._jnetwork.saveToString()
        reparsed = minidom.parseString(nt)
        with open(path, 'w') as fh:
            fh.write(reparsed.toprettyxml(indent="  "))

    def is_trained(self):
        return bayespy.network.is_trained(self._jnetwork)

    def train(self):
        """
        Train a model on data provided in the constructor
        """
        learning = bayesServerParams().ParameterLearning(self._jnetwork, self._inference_factory.get_inference_factory())
        learning_options = bayesServerParams().ParameterLearningOptions()

        data_reader_command = self._data_store.create_data_reader_command()

        reader_options = bayesServer().data.ReaderOptions()

        variable_references = list(bayespy.network.create_variable_references(self._jnetwork, self._data))

        evidence_reader_command = bayesServer().data.DefaultEvidenceReaderCommand(data_reader_command,
                                            jp.java.util.Arrays.asList(variable_references), reader_options)

        self._logger.info("Training model...")
        result = learning.learn(evidence_reader_command, learning_options)
        self._logger.info("Finished training model")

        return {'Converged': result.getConverged(), 'Loglikelihood': result.getLogLikelihood().floatValue(),
                    'IterationCount': result.getIterationCount(), 'CaseCount': result.getCaseCount(),
                    'WeightedCaseCount': result.getWeightedCaseCount(), 'UnweightedCaseCount':  result.getUnweightedCaseCount(),
                    'BIC': result.getBIC().floatValue()}

    def batch_query(self, queries=[QueryStatistics()], append_to_df=True, variable_references=[]):
        bq = BatchQuery(self._jnetwork, self._data_store, self._logger)
        return bq.query(queries, append_to_df=append_to_df, variable_references=variable_references)<commit_msg>Forcing spawning in Linux<commit_after># -*- coding: utf-8 -*-
"""
Created on Tue Apr 19 16:31:18 2016

@author: imorgan.admin
"""

from bayespy.jni import bayesServerInference

import bayespy.network
import pandas as pd
from bayespy.jni import bayesServer
from bayespy.jni import bayesServerParams
from bayespy.jni import jp
import numpy as np
import logging
from pathos.multiprocessing import ProcessPool
import multiprocess.context as ctx
import pathos.multiprocessing as mp

from typing import List

class QueryOutput:
    def __init__(self, continuous, discrete):
        self.continuous = continuous
        self.discrete = discrete

class InferenceEngine:

    _inference_factory = None

    def __init__(self, network):
        self._network = network

    @staticmethod
    def get_inference_factory():
        if InferenceEngine._inference_factory is None:
            InferenceEngine._inference_factory = bayesServerInference().RelevanceTreeInferenceFactory()
        return InferenceEngine._inference_factory

    def create(self, loglikelihood=False, conflict=False, retract=False):
        query_options = self.get_inference_factory().createQueryOptions()
        query_output = self.get_inference_factory().createQueryOutput()
        inference_engine = self.get_inference_factory().createInferenceEngine(self._network)

        query_options.setLogLikelihood(loglikelihood)
        query_options.setConflict(conflict)

        if retract:
            query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)

        return inference_engine, query_options, query_output

class Query:

    def __init__(self, jnetwork, inference):
        self._factory = bayesServerInference().RelevanceTreeInferenceFactory()
        self._queryOptions = self._factory.createQueryOptions()
        self._queryOutput = self._factory.createQueryOutput()
        self._queryDistributions = inference.getQueryDistributions()
        self._inference = inference
        self._jnetwork = jnetwork

    def execute(self, variables=None):
        """
        Query a number of variables (if none, then query all variables in the network)
        :param variables: a list of variables, or none
        :return: a QueryOutput object with separate continuous/ discrete dataframes
        """
        distributions = {}
        if variables is None:
            variables = self._jnetwork.getVariables()

        for v in variables:
            if bayespy.network.is_variable_discrete(v):
                table = bayesServer().Table(v)
            else:
                table = bayesServer().CLGaussian(v)

            self._queryDistributions.add(bayesServerInference().QueryDistribution(table))
            distributions.update({ v.getName() : table})

        self._inference.query(self._queryOptions, self._queryOutput)

        states= []
        d_variables = []
        values = []
        c_variables = []
        mean = []
        variance = []
        for v in variables:
            dist =  distributions[v.getName()]
            if bayespy.network.is_variable_discrete(v):
                for state in v.getStates():
                    states.append(state.getName())
                    values.append(float(dist.get([state])))
                    d_variables.append(v.getName())
            else:
                variance.append(float(dist.getVariance(v)))
                mean.append(float(dist.getMean(v)))
                c_variables.append(v.getName())

        return QueryOutput(pd.DataFrame({ 'variable': c_variables, 'mean': mean, 'variance': variance}), pd.DataFrame({ 'variable' : d_variables, 'state': states, 'value' : values}))

class Evidence:

    def __init__(self, jnetwork, inference):
        self._jnetwork = jnetwork
        self._inference = inference
        self._evidence = inference.getEvidence()
        self._evidence.clear()
        self._variables = jnetwork.getVariables()

    def apply(self, evidence):
        """
        Apply evidence to a network
        :param evidence: if a discrete variable, then a list of strings generated by network.Discrete.tostring(),
        if continuous, then a list of tuples with (VariableName, ContinuousValue)
        :return: Nothing
        """
        for value in evidence:
            if not isinstance(value, tuple):
                value = str(value)
                node, state = value.split(bayespy.network.STATE_DELIMITER)
                v = self._variables.get(node)
                if v is None:
                    raise ValueError("Node {} does not exist".format(node))
                if bayespy.network.is_variable_discrete(v):
                    st = v.getStates().get(state)
                    if st is None:
                        raise ValueError("State {} does not exist in variable {}".format(state, node))

                    self._evidence.setState(st)
                else:
                    raise ValueError("Can not find state on a continuous variable")
            else:
                v = self._variables.get(value[0])
                if not bayespy.network.is_variable_continuous(v):
                    raise ValueError("Variable is not continuous")

                self._evidence.set(v, jp.java.lang.Double(value[1]))

        return self._evidence

class QueryStatistics:

    def __init__(self, calc_loglikelihood=True, calc_conflict=False, loglikelihood_column='loglikelihood', conflict_column='conflict'):
        self._calc_loglikelihood = calc_loglikelihood
        self._calc_conflict = calc_conflict
        self._loglikelihood_column = loglikelihood_column
        self._conflict_column = conflict_column

    def setup(self, network, inference_engine, query_options):
        query_options.setLogLikelihood(self._calc_loglikelihood)
        query_options.setConflict(self._calc_conflict)

    def results(self, inference_engine, query_output):
        result = {}
        if self._calc_loglikelihood:
            result.update({ self._loglikelihood_column: query_output.getLogLikelihood().floatValue() })

        if self._calc_conflict:
            result.update({ self._conflict_column: query_output.getConflict().floatValue()})

        return result

# seems like a better name than QueryStatistics, so just having this here.
class QueryModelStatistics(QueryStatistics):
    def __init__(self, calc_loglikelihood=True, calc_conflict=False, loglikelihood_column='loglikelihood', conflict_column='conflict'):
        super().__init__(calc_loglikelihood, calc_conflict, loglikelihood_column, conflict_column)

class QueryMostLikelyState:

    def __init__(self, target_variable_name, output_dtype="object", suffix="_maxlikelihood"):
        self._target_variable_name = target_variable_name
        self._distribution = None
        self._output_dtype = output_dtype
        self._suffix = suffix

    def setup(self, network, inference_engine, query_options):
        distribution = None

        self._variable = bayespy.network.get_variable(network, self._target_variable_name)

        if bayespy.network.is_variable_discrete(self._variable):
            distribution = bayesServer().Table(self._variable)

        if distribution is None:
            raise ValueError("{} needs to be discrete in QueryMostLikelyState".format(self._target_variable_name))

        query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)
        qd = bayesServerInference().QueryDistribution(distribution)

        self._distribution = distribution
        inference_engine.getQueryDistributions().add(qd)

    def results(self, inference_engine, query_output):
        states = {}
        for state in self._variable.getStates():
            states.update({ state.getName() : self._distribution.get([state])})

        # get the most likely state
        max_state = max(states.keys(), key=(lambda key: states[key]))
        max_state_name = bayespy.data.DataFrame.cast2(self._output_dtype, max_state)

        return {self._target_variable_name + self._suffix: max_state_name}

class QueryLogLikelihood:
    def __init__(self, variable_names, column_name: str='_loglikelihood'):
        if isinstance(variable_names, str):
            variable_names = [variable_names]

        self._variable_names = variable_names
        self._distribution = None
        self._query_distribution = None

        self._column_name = column_name

    def setup(self, network, inference_engine, query_options):
        variables = [bayespy.network.get_variable(network, n) for n in self._variable_names]
        if len(variables) == 1:
            self._distribution = bayesServer().CLGaussian(variables[0])
        else:
            self._distribution = bayesServer().CLGaussian(variables)

        query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)
        qd = bayesServerInference().QueryDistribution(self._distribution)
        qd.setQueryLogLikelihood(True)
        self._query_distribution = qd
        inference_engine.getQueryDistributions().add(qd)

    def results(self, inference_engine, query_output):
        result = {}
        ll = self._query_distribution.getLogLikelihood()
        value = ll.floatValue() if ll is not None else np.nan
        result.update({":".join(self._variable_names) + self._column_name: value})
        return result

class QueryMeanVariance:
    def __init__(self, variable_name, retract_evidence=True, result_mean_suffix='_mean', result_variance_suffix='_variance'):
        self._variable_name = variable_name

        self._result_mean_suffix = result_mean_suffix
        self._result_variance_suffix = result_variance_suffix
        self._retract_evidence = retract_evidence

    def setup(self, network, inference_engine, query_options):
        self._variable = bayespy.network.get_variable(network, self._variable_name)

        if not bayespy.network.is_variable_continuous(self._variable):
            raise ValueError("{} needs to be continuous.".format(self._variable_name))

        self._query = bayesServer().CLGaussian(self._variable)

        if self._retract_evidence:
            query_options.setQueryEvidenceMode(bayesServerInference().QueryEvidenceMode.RETRACT_QUERY_EVIDENCE)

        inference_engine.getQueryDistributions().add(bayesServerInference().QueryDistribution(self._query))

    def results(self, inference_engine, query_output):
        return {self._variable_name + self._result_mean_suffix: self._query.getMean(self._variable),
                self._variable_name + self._result_variance_suffix: self._query.getVariance(self._variable)}

def _batch_query(df: pd.DataFrame, connection_string: str, network: str, table_name: str,
                 variable_references: List[str],
                 queries, logger, i):

    bayespy.jni.attach(logger)
    data_reader = bayesServer().data.DatabaseDataReaderCommand(
        connection_string,
        "select * from {} where ix in ({})".format(table_name,
                                                   ",".join(str(i) for i in df.index.tolist()))).executeReader()

    network = bayespy.network.create_network_from_string(network)
    reader_options = bayesServer().data.ReaderOptions("ix")
    variable_refs = list(bayespy.network.create_variable_references(network, df,
                                                                    variable_references=variable_references))

    reader = bayesServer().data.DefaultEvidenceReader(data_reader, jp.java.util.Arrays.asList(variable_refs),
                                                   reader_options)

    factory = InferenceEngine(network)
    (inference_engine, query_options, query_output) = factory.create()

    for query in queries:
        query.setup(network, inference_engine, query_options)

    results = []
    try:
        while reader.read(inference_engine.getEvidence(), bayesServer().data.DefaultReadOptions(True)):
            result = {}

            try:
                inference_engine.query(query_options, query_output)
            except BaseException as e:
                logger.error(e)
                # inference_engine.getEvidence().clear()
                # continue

            for query in queries:
                result = {**result, **query.results(inference_engine, query_output)}

            inference_engine.getEvidence().clear()
            result.update({'caseid': int(reader.getReadInfo().getCaseId().toString())})

            results.append(result)

            if i % 500 == 0:
                logger.info("Queried case {}".format(i))

            i += 1
    finally:
        reader.close()
        #bayespy.jni.detach()
    return results

class BatchQuery:

    def __init__(self, network, datastore, logger: logging.Logger):

        self._logger = logger
        self._datastore = datastore
        # serialise the network as a string.
        self._network = network.saveToString()

    def _calc_num_threads(self, df_size: int, query_size: int) -> int:
        num_queries = df_size * query_size

        if mp.cpu_count() == 1:
            max = 1
        else:
            max = mp.cpu_count() - 1

        calc = int(num_queries / 5000)
        if calc > max:
            r = max
        elif calc <= 1:
            if num_queries > 1000:
                r = 2
            else:
                r = 1
        else:
            r = calc

        return r

    def query(self, queries=[QueryStatistics()], append_to_df=True, variable_references=[]):

        if not hasattr(queries, "__getitem__"):
            queries = [queries]

        nt = self._network
        logger = self._logger
        conn = self._datastore.get_connection()
        table = self._datastore.table
        processes = self._calc_num_threads(len(self._datastore.data), len(queries))

        self._logger.info("Using {} processes to query {} rows".format(processes, len(self._datastore.data)))

        if processes == 1:
            pdf = pd.DataFrame(_batch_query(self._datastore.data, conn, nt, table,
                         variable_references, queries,
                         logger, 0))
        else:
            # bit nasty, but the only way I could get jpype to stop hanging in Linux.
            ctx._force_start_method('spawn')

            with ProcessPool(nodes=processes) as pool:
                pdf = pd.DataFrame()
                for result_set in pool.map(lambda df: _batch_query(df, conn, nt, table,
                                                                 variable_references, queries,
                                                                logger, 0), np.array_split(self._datastore.data, processes)):
                    pdf = pdf.append(pd.DataFrame(result_set))

        df = pdf.set_index('caseid')

        if append_to_df:
            return self._datastore.data.join(df)
        else:
            return df

class NetworkModel:

    def __init__(self, network, data_store, logger):
        self._jnetwork = network
        self._inference_factory = InferenceEngine(network)
        self._data_store = data_store
        self._logger = logger
        self._data = data_store.data

    def get_network(self):
        return self._jnetwork

    def save(self, path):
        from xml.dom import minidom
        nt = self._jnetwork.saveToString()
        reparsed = minidom.parseString(nt)
        with open(path, 'w') as fh:
            fh.write(reparsed.toprettyxml(indent="  "))

    def is_trained(self):
        return bayespy.network.is_trained(self._jnetwork)

    def train(self):
        """
        Train a model on data provided in the constructor
        """
        learning = bayesServerParams().ParameterLearning(self._jnetwork, self._inference_factory.get_inference_factory())
        learning_options = bayesServerParams().ParameterLearningOptions()

        data_reader_command = self._data_store.create_data_reader_command()

        reader_options = bayesServer().data.ReaderOptions()

        variable_references = list(bayespy.network.create_variable_references(self._jnetwork, self._data))

        evidence_reader_command = bayesServer().data.DefaultEvidenceReaderCommand(data_reader_command,
                                            jp.java.util.Arrays.asList(variable_references), reader_options)

        self._logger.info("Training model...")
        result = learning.learn(evidence_reader_command, learning_options)
        self._logger.info("Finished training model")

        return {'Converged': result.getConverged(), 'Loglikelihood': result.getLogLikelihood().floatValue(),
                    'IterationCount': result.getIterationCount(), 'CaseCount': result.getCaseCount(),
                    'WeightedCaseCount': result.getWeightedCaseCount(), 'UnweightedCaseCount':  result.getUnweightedCaseCount(),
                    'BIC': result.getBIC().floatValue()}

    def batch_query(self, queries=[QueryStatistics()], append_to_df=True, variable_references=[]):
        bq = BatchQuery(self._jnetwork, self._data_store, self._logger)
        return bq.query(queries, append_to_df=append_to_df, variable_references=variable_references)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (C) 2004-2007 Edgewall Software
# Copyright (C) 2004-2007 Christopher Lenz <cmlenz@gmx.de>
# Copyright (C) 2005 Matthew Good <trac@matt-good.net>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>
#         Matthew Good <trac@matt-good.net>

import os
import pkg_resources
import urllib

from mod_python import apache
try:
    from mod_python import version
except ImportError:
    version = "< 3.2"

from trac import __version__ as VERSION
from trac.web.main import dispatch_request
from trac.web.wsgi import WSGIGateway, _ErrorsWrapper

class InputWrapper(object):

    def __init__(self, req):
        self.req = req

    def close(self):
        pass

    def read(self, size=-1):
        return self.req.read(size)

    def readline(self, size=-1):
        return self.req.readline(size)

    def readlines(self, hint=-1):
        return self.req.readlines(hint)

class ModPythonGateway(WSGIGateway):

    wsgi_multithread = apache.mpm_query(apache.AP_MPMQ_IS_THREADED) > 0
    wsgi_multiprocess = apache.mpm_query(apache.AP_MPMQ_IS_FORKED) > 0

    def __init__(self, req, options):
        environ = {}
        environ.update(apache.build_cgi_env(req))

        environ['trac.web.frontend'] = 'mod_python'
        environ['trac.web.version'] = version

        if 'TracEnv' in options:
            environ['trac.env_path'] = options['TracEnv']
        if 'TracEnvParentDir' in options:
            environ['trac.env_parent_dir'] = options['TracEnvParentDir']
        if 'TracEnvIndexTemplate' in options:
            environ['trac.env_index_template'] = options['TracEnvIndexTemplate']
        if 'TracTemplateVars' in options:
            environ['trac.template_vars'] = options['TracTemplateVars']
        if 'TracLocale' in options:
            environ['trac.locale'] = options['TracLocale']

        if 'TracUriRoot' in options:
            # Special handling of SCRIPT_NAME/PATH_INFO for mod_python, which
            # tends to get confused for whatever reason
            root_uri = options['TracUriRoot'].rstrip('/')
            request_uri = environ['REQUEST_URI'].split('?', 1)[0]
            if not request_uri.startswith(root_uri):
                raise ValueError('TracUriRoot set to %s but request URL '
                                 'is %s' % (root_uri, request_uri))
            environ['SCRIPT_NAME'] = root_uri
            environ['PATH_INFO'] = urllib.unquote(request_uri[len(root_uri):])

        egg_cache = req.subprocess_env.get('PYTHON_EGG_CACHE')
        if egg_cache:
            os.environ['PYTHON_EGG_CACHE'] = egg_cache

        WSGIGateway.__init__(self, environ, InputWrapper(req),
                             _ErrorsWrapper(lambda x: req.log_error(x)))
        self.req = req

    def _send_headers(self):
        assert self.headers_set, 'Response not started'

        if not self.headers_sent:
            status, headers = self.headers_sent = self.headers_set
            self.req.status = int(status[:3])
            for name, value in headers:
                if name.lower() == 'content-length':
                    self.req.set_content_length(int(value))
                elif name.lower() == 'content-type':
                    self.req.content_type = value
                else:
                    self.req.headers_out.add(name, value)

    def _sendfile(self, fileobj):
        self._send_headers()
        try:
            self.req.sendfile(fileobj.name)
        except IOError, e:
            if 'client closed connection' not in str(e):
                raise

    def _write(self, data):
        self._send_headers()
        try:
            self.req.write(data)
        except IOError, e:
            if 'client closed connection' not in str(e):
                raise

def handler(req):
    pkg_resources.require('Trac==%s' % VERSION)
    gateway = ModPythonGateway(req, req.get_options())
    gateway.run(dispatch_request)
    return apache.OK
<commit_msg>Move the dispatch_request import so PYTHON_EGG_CACHE will be set first. This  is important if Genshi (which is imported by trac.web.main) is installed as a zipped egg.<commit_after># -*- coding: utf-8 -*-
#
# Copyright (C) 2004-2007 Edgewall Software
# Copyright (C) 2004-2007 Christopher Lenz <cmlenz@gmx.de>
# Copyright (C) 2005 Matthew Good <trac@matt-good.net>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Christopher Lenz <cmlenz@gmx.de>
#         Matthew Good <trac@matt-good.net>

import os
import pkg_resources
import urllib

from mod_python import apache
try:
    from mod_python import version
except ImportError:
    version = "< 3.2"

from trac import __version__ as VERSION
from trac.web.wsgi import WSGIGateway, _ErrorsWrapper

class InputWrapper(object):

    def __init__(self, req):
        self.req = req

    def close(self):
        pass

    def read(self, size=-1):
        return self.req.read(size)

    def readline(self, size=-1):
        return self.req.readline(size)

    def readlines(self, hint=-1):
        return self.req.readlines(hint)

class ModPythonGateway(WSGIGateway):

    wsgi_multithread = apache.mpm_query(apache.AP_MPMQ_IS_THREADED) > 0
    wsgi_multiprocess = apache.mpm_query(apache.AP_MPMQ_IS_FORKED) > 0

    def __init__(self, req, options):
        environ = {}
        environ.update(apache.build_cgi_env(req))

        environ['trac.web.frontend'] = 'mod_python'
        environ['trac.web.version'] = version

        if 'TracEnv' in options:
            environ['trac.env_path'] = options['TracEnv']
        if 'TracEnvParentDir' in options:
            environ['trac.env_parent_dir'] = options['TracEnvParentDir']
        if 'TracEnvIndexTemplate' in options:
            environ['trac.env_index_template'] = options['TracEnvIndexTemplate']
        if 'TracTemplateVars' in options:
            environ['trac.template_vars'] = options['TracTemplateVars']
        if 'TracLocale' in options:
            environ['trac.locale'] = options['TracLocale']

        if 'TracUriRoot' in options:
            # Special handling of SCRIPT_NAME/PATH_INFO for mod_python, which
            # tends to get confused for whatever reason
            root_uri = options['TracUriRoot'].rstrip('/')
            request_uri = environ['REQUEST_URI'].split('?', 1)[0]
            if not request_uri.startswith(root_uri):
                raise ValueError('TracUriRoot set to %s but request URL '
                                 'is %s' % (root_uri, request_uri))
            environ['SCRIPT_NAME'] = root_uri
            environ['PATH_INFO'] = urllib.unquote(request_uri[len(root_uri):])

        egg_cache = req.subprocess_env.get('PYTHON_EGG_CACHE')
        if egg_cache:
            os.environ['PYTHON_EGG_CACHE'] = egg_cache

        WSGIGateway.__init__(self, environ, InputWrapper(req),
                             _ErrorsWrapper(lambda x: req.log_error(x)))
        self.req = req

    def _send_headers(self):
        assert self.headers_set, 'Response not started'

        if not self.headers_sent:
            status, headers = self.headers_sent = self.headers_set
            self.req.status = int(status[:3])
            for name, value in headers:
                if name.lower() == 'content-length':
                    self.req.set_content_length(int(value))
                elif name.lower() == 'content-type':
                    self.req.content_type = value
                else:
                    self.req.headers_out.add(name, value)

    def _sendfile(self, fileobj):
        self._send_headers()
        try:
            self.req.sendfile(fileobj.name)
        except IOError, e:
            if 'client closed connection' not in str(e):
                raise

    def _write(self, data):
        self._send_headers()
        try:
            self.req.write(data)
        except IOError, e:
            if 'client closed connection' not in str(e):
                raise

def handler(req):
    pkg_resources.require('Trac==%s' % VERSION)
    gateway = ModPythonGateway(req, req.get_options())
    from trac.web.main import dispatch_request
    gateway.run(dispatch_request)
    return apache.OK
<|endoftext|>


<commit_before># -*- coding: iso8859-1 -*-
#
# Copyright (C) 2004, 2005 Edgewall Software
# Copyright (C) 2004, 2005 Christopher Lenz <cmlenz@gmx.de>
#
# Trac is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the
# License, or (at your option) any later version.
#
# Trac 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
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
# Author: Christopher Lenz <cmlenz@gmx.de>

import locale
import mimetypes
import os
import re
import threading

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from mod_python import apache, util

from trac.env import open_environment
from trac.util import TracError, enum, href_join, http_date, rstrip
from trac.web.main import Request, RequestDone, dispatch_request, \
                          send_pretty_error

class ModPythonRequest(Request):

    idx_location = None

    def __init__(self, req, options):
        Request.__init__(self)
        self.req = req

        self.method = self.req.method
        self.server_name = self.req.server.server_hostname
        self.server_port = self.req.connection.local_addr[1]
        self.remote_addr = self.req.connection.remote_ip
        self.remote_user = self.req.user
        self.scheme = 'http'
        if self.req.subprocess_env.get('HTTPS') in ('on', '1') \
                or self.server_port == 443:
            self.scheme = 'https'
        if self.req.headers_in.has_key('Cookie'):
            self.incookie.load(self.req.headers_in['Cookie'])
        self.args = FieldStorageWrapper(self.req)

        # The root uri sometimes has to be explicitly specified because apache
        # sometimes get req.path_info wrong if many <alias> and <location> directives
        # are used.
        if options.has_key('TracUriRoot'):
            root_uri = rstrip(options['TracUriRoot'], '/')
            if self.req.uri[:len(root_uri)] != root_uri:
                raise ValueError('TracRootUri set to "%s" but req.uri starts with "%s"' %
                                 (root_uri, self.req.uri[:len(root_uri)]))
            self.path_info = self.req.uri[len(root_uri):]
        else:
            self.path_info = self.req.path_info

        if len(self.path_info):
            self.idx_location = self.req.uri[:-len(self.path_info)]
        else:
            self.idx_location = self.req.uri

        if options.has_key('TracEnvParentDir') and self.path_info:
            # We have to remove one path element from path_info when we're
            # using TracEnvParentDir
            self.path_info = re.sub('/[^/]+', '', self.path_info, 1)

        if len(self.path_info):
            self.cgi_location = self.req.uri[:-len(self.path_info)] or '/'
        else:
            self.cgi_location = self.req.uri

    def read(self, len):
        return self.req.read(len)

    def write(self, data):
        self.req.write(data)

    def get_header(self, name):
        return self.req.headers_in.get(name)

    def send_response(self, code):
        self.req.status = code

    def send_header(self, name, value):
        if name.lower() == 'content-type':
            self.req.content_type = value
        elif name.lower() == 'content-length':
            self.req.set_content_length(int(value))
        else:
            self.req.headers_out.add(name, str(value))

    def end_headers(self):
        pass

    def send_file(self, path, mimetype=None):
        stat = os.stat(path)
        last_modified = http_date(stat.st_mtime)
        if last_modified == self.req.headers_in.get('If-Modified-Since'):
            self.send_response(304)
            raise RequestDone

        self.req.status = 200
        if not mimetype:
            mimetype = mimetypes.guess_type(path)[0]
        self.req.content_type = mimetype
        self.req.set_content_length(stat.st_size)
        self.req.headers_out.add('Last-Modified', http_date(stat.st_mtime))

        self.req.sendfile(path)
        raise RequestDone

class FieldStorageWrapper(util.FieldStorage):
    """
    mod_python FieldStorage wrapper that improves compatibility with the other
    front-ends.
    """

    def __init__(self, req):
        """
        The mod_python FieldStorage implementation, unlike cgi.py, always
        includes GET parameters, even if they are also defined in the body of
        a POST request. We work around this to provide the behaviour of cgi.py
        here.
        """
        class RequestWrapper(object):
            def __init__(self, req):
                self.req = req
                self.args = ''
            def __getattr__(self, name):
                return getattr(self.req, name)
        util.FieldStorage.__init__(self, RequestWrapper(req), keep_blank_values=1)

        # Populate FieldStorage with the original query string parameters, if
        # they aren't already defined through the request body
        if req.args:
            qsargs = []
            for pair in util.parse_qsl(req.args, 1):
                if self.has_key(pair[0]):
                    continue
                qsargs.append(util.Field(pair[0], StringIO(pair[1]),
                                         "text/plain", {}, None, {}))
            self.list += qsargs

    def get(self, key, default=None):
        return util.FieldStorage.get(self, key, default)

    def __setitem__(self, key, value):
        if not key in self:
            self.list.append(util.Field(key, StringIO(value), 'text/plain',
                             {}, None, {}))

def send_project_index(req, mpr, dir, options):
    from trac.web.clearsilver import HDFWrapper

    if 'TracEnvIndexTemplate' in options:
        # Custom project listing template configured
        tmpl_path, template = os.path.split(options['TracEnvIndexTemplate'])

        from trac.config import default_dir
        mpr.hdf = HDFWrapper(loadpaths=[default_dir('templates'), tmpl_path])

        tmpl_vars = {}
        if 'TracTemplateVars' in options:
            pairs = options['TracTemplateVars'].split(',')
            for pair in pairs:
                key,val = pair.split('=')
                mpr.hdf[key] = val

    else:
        # Use the default project listing template
        mpr.hdf = HDFWrapper()
        template = mpr.hdf.parse("""
<title>Available Projects</title>

Available Projects

<?cs
 each:project = projects ?>
<a href="<?cs
  var:project.href ?>"><?cs var:project.name ?>
<?cs
 /each ?>

""")

    try:
        projects = []
        for idx, project in enum(os.listdir(dir)):
            env_path = os.path.join(dir, project)
            if not os.path.isdir(env_path):
                continue
            try:
                env = open_environment(env_path)
                projects.append({
                    'name': env.config.get('project', 'name'),
                    'description': env.config.get('project', 'descr'),
                    'href': href_join(mpr.idx_location, project)
                })
            except TracError, e:
                req.log_error('Error opening environment at %s: %s'
                              % (env_path, e))
        projects.sort(lambda x, y: cmp(x['name'], y['name']))
        mpr.hdf['projects'] = projects
        mpr.display(template, response=200)
    except RequestDone:
        pass

env_cache = {}
env_cache_lock = threading.Lock()

def get_environment(req, mpr, options):
    global env_cache, env_cache_lock

    if options.has_key('TracEnv'):
        env_path = options['TracEnv']
    elif options.has_key('TracEnvParentDir'):
        env_parent_dir = options['TracEnvParentDir']
        env_name = mpr.cgi_location.split('/')[-1]
        env_path = os.path.join(env_parent_dir, env_name)
        if len(env_name) == 0 or not os.path.exists(env_path):
            send_project_index(req, mpr, env_parent_dir, options)
            return None
    else:
        raise TracError, \
              'Missing PythonOption "TracEnv" or "TracEnvParentDir". Trac ' \
              'requires one of these options to locate the Trac environment(s).'

    env = None
    try:
        env_cache_lock.acquire()
        if not env_path in env_cache:
            env_cache[env_path] = open_environment(env_path)
        env = env_cache[env_path]
    finally:
        env_cache_lock.release()
    return env

def handler(req):
    options = req.get_options()
    if options.has_key('TracLocale'):
        locale.setlocale(locale.LC_ALL, options['TracLocale'])
    else:
        locale.setlocale(locale.LC_ALL, '')

    mpr = ModPythonRequest(req, options)
    env = get_environment(req, mpr, options)
    if not env:
        return apache.OK

    req.content_type = 'text/html'
    try:
        dispatch_request(mpr.path_info, mpr, env)
    except Exception, e:
        send_pretty_error(e, env, mpr)
    return apache.OK
<commit_msg>mod_python handler: Only release lock if it has been acquired.<commit_after># -*- coding: iso8859-1 -*-
#
# Copyright (C) 2004, 2005 Edgewall Software
# Copyright (C) 2004, 2005 Christopher Lenz <cmlenz@gmx.de>
#
# Trac is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the
# License, or (at your option) any later version.
#
# Trac 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
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
#
# Author: Christopher Lenz <cmlenz@gmx.de>

import locale
import mimetypes
import os
import re
import threading

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from mod_python import apache, util

from trac.env import open_environment
from trac.util import TracError, enum, href_join, http_date, rstrip
from trac.web.main import Request, RequestDone, dispatch_request, \
                          send_pretty_error

class ModPythonRequest(Request):

    idx_location = None

    def __init__(self, req, options):
        Request.__init__(self)
        self.req = req

        self.method = self.req.method
        self.server_name = self.req.server.server_hostname
        self.server_port = self.req.connection.local_addr[1]
        self.remote_addr = self.req.connection.remote_ip
        self.remote_user = self.req.user
        self.scheme = 'http'
        if self.req.subprocess_env.get('HTTPS') in ('on', '1') \
                or self.server_port == 443:
            self.scheme = 'https'
        if self.req.headers_in.has_key('Cookie'):
            self.incookie.load(self.req.headers_in['Cookie'])
        self.args = FieldStorageWrapper(self.req)

        # The root uri sometimes has to be explicitly specified because apache
        # sometimes get req.path_info wrong if many <alias> and <location> directives
        # are used.
        if options.has_key('TracUriRoot'):
            root_uri = rstrip(options['TracUriRoot'], '/')
            if self.req.uri[:len(root_uri)] != root_uri:
                raise ValueError('TracRootUri set to "%s" but req.uri starts with "%s"' %
                                 (root_uri, self.req.uri[:len(root_uri)]))
            self.path_info = self.req.uri[len(root_uri):]
        else:
            self.path_info = self.req.path_info

        if len(self.path_info):
            self.idx_location = self.req.uri[:-len(self.path_info)]
        else:
            self.idx_location = self.req.uri

        if options.has_key('TracEnvParentDir') and self.path_info:
            # We have to remove one path element from path_info when we're
            # using TracEnvParentDir
            self.path_info = re.sub('/[^/]+', '', self.path_info, 1)

        if len(self.path_info):
            self.cgi_location = self.req.uri[:-len(self.path_info)] or '/'
        else:
            self.cgi_location = self.req.uri

    def read(self, len):
        return self.req.read(len)

    def write(self, data):
        self.req.write(data)

    def get_header(self, name):
        return self.req.headers_in.get(name)

    def send_response(self, code):
        self.req.status = code

    def send_header(self, name, value):
        if name.lower() == 'content-type':
            self.req.content_type = value
        elif name.lower() == 'content-length':
            self.req.set_content_length(int(value))
        else:
            self.req.headers_out.add(name, str(value))

    def end_headers(self):
        pass

    def send_file(self, path, mimetype=None):
        stat = os.stat(path)
        last_modified = http_date(stat.st_mtime)
        if last_modified == self.req.headers_in.get('If-Modified-Since'):
            self.send_response(304)
            raise RequestDone

        self.req.status = 200
        if not mimetype:
            mimetype = mimetypes.guess_type(path)[0]
        self.req.content_type = mimetype
        self.req.set_content_length(stat.st_size)
        self.req.headers_out.add('Last-Modified', http_date(stat.st_mtime))

        self.req.sendfile(path)
        raise RequestDone

class FieldStorageWrapper(util.FieldStorage):
    """
    mod_python FieldStorage wrapper that improves compatibility with the other
    front-ends.
    """

    def __init__(self, req):
        """
        The mod_python FieldStorage implementation, unlike cgi.py, always
        includes GET parameters, even if they are also defined in the body of
        a POST request. We work around this to provide the behaviour of cgi.py
        here.
        """
        class RequestWrapper(object):
            def __init__(self, req):
                self.req = req
                self.args = ''
            def __getattr__(self, name):
                return getattr(self.req, name)
        util.FieldStorage.__init__(self, RequestWrapper(req), keep_blank_values=1)

        # Populate FieldStorage with the original query string parameters, if
        # they aren't already defined through the request body
        if req.args:
            qsargs = []
            for pair in util.parse_qsl(req.args, 1):
                if self.has_key(pair[0]):
                    continue
                qsargs.append(util.Field(pair[0], StringIO(pair[1]),
                                         "text/plain", {}, None, {}))
            self.list += qsargs

    def get(self, key, default=None):
        return util.FieldStorage.get(self, key, default)

    def __setitem__(self, key, value):
        if not key in self:
            self.list.append(util.Field(key, StringIO(value), 'text/plain',
                             {}, None, {}))

def send_project_index(req, mpr, dir, options):
    from trac.web.clearsilver import HDFWrapper

    if 'TracEnvIndexTemplate' in options:
        # Custom project listing template configured
        tmpl_path, template = os.path.split(options['TracEnvIndexTemplate'])

        from trac.config import default_dir
        mpr.hdf = HDFWrapper(loadpaths=[default_dir('templates'), tmpl_path])

        tmpl_vars = {}
        if 'TracTemplateVars' in options:
            pairs = options['TracTemplateVars'].split(',')
            for pair in pairs:
                key,val = pair.split('=')
                mpr.hdf[key] = val

    else:
        # Use the default project listing template
        mpr.hdf = HDFWrapper()
        template = mpr.hdf.parse("""
<title>Available Projects</title>

Available Projects

<?cs
 each:project = projects ?>
<a href="<?cs
  var:project.href ?>"><?cs var:project.name ?>
<?cs
 /each ?>

""")

    try:
        projects = []
        for idx, project in enum(os.listdir(dir)):
            env_path = os.path.join(dir, project)
            if not os.path.isdir(env_path):
                continue
            try:
                env = open_environment(env_path)
                projects.append({
                    'name': env.config.get('project', 'name'),
                    'description': env.config.get('project', 'descr'),
                    'href': href_join(mpr.idx_location, project)
                })
            except TracError, e:
                req.log_error('Error opening environment at %s: %s'
                              % (env_path, e))
        projects.sort(lambda x, y: cmp(x['name'], y['name']))
        mpr.hdf['projects'] = projects
        mpr.display(template, response=200)
    except RequestDone:
        pass

env_cache = {}
env_cache_lock = threading.Lock()

def get_environment(req, mpr, options):
    global env_cache, env_cache_lock

    if options.has_key('TracEnv'):
        env_path = options['TracEnv']
    elif options.has_key('TracEnvParentDir'):
        env_parent_dir = options['TracEnvParentDir']
        env_name = mpr.cgi_location.split('/')[-1]
        env_path = os.path.join(env_parent_dir, env_name)
        if len(env_name) == 0 or not os.path.exists(env_path):
            send_project_index(req, mpr, env_parent_dir, options)
            return None
    else:
        raise TracError, \
              'Missing PythonOption "TracEnv" or "TracEnvParentDir". Trac ' \
              'requires one of these options to locate the Trac environment(s).'

    env = None
    env_cache_lock.acquire()
    try:
        if not env_path in env_cache:
            env_cache[env_path] = open_environment(env_path)
        env = env_cache[env_path]
    finally:
        env_cache_lock.release()
    return env

def handler(req):
    options = req.get_options()
    if options.has_key('TracLocale'):
        locale.setlocale(locale.LC_ALL, options['TracLocale'])
    else:
        locale.setlocale(locale.LC_ALL, '')

    mpr = ModPythonRequest(req, options)
    env = get_environment(req, mpr, options)
    if not env:
        return apache.OK

    req.content_type = 'text/html'
    try:
        dispatch_request(mpr.path_info, mpr, env)
    except Exception, e:
        send_pretty_error(e, env, mpr)
    return apache.OK
<|endoftext|>


<commit_before>from piptools.datastructures import SpecSet, Spec

def print_specset(specset, round):
    print('After round #%s:' % (round,))
    for spec in specset:
        print('  - %s' % (spec.description(),))

class Resolver(object):
    def __init__(self, spec_set, package_manager=None):
        """This class resolves a given SpecSet by querying the given
        PackageManager.
        """
        # TODO: If no specific PackageManager instance is used, use the
        # default package manager (locally cached PyPI).
        self.spec_set = spec_set
        #if package_manager is None:
        #    from piptools.package_manager import PackageManager
        #    package_manager = PackageManager(...)
        self.pkgmgr = package_manager

    def resolve_one_round(self):
        """Resolves one level of the current spec set, by finding best matches
        for the current spec set in the package manager and returning all
        (new) requirements for those packages.

        Returns whether the spec set was changed significantly by this round.
        """
        new_deps = self.find_new_dependencies()
        self.spec_set.add_specs(new_deps)
        return len(new_deps) > 0

    def resolve(self, max_rounds=12):
        """Resolves the spec set one round at a time, until the set does not
        change significantly anymore.  Protects against infinite loops by
        breaking out after a max number rounds.
        """
        round = 0
        while True:
            round += 1
            if round > max_rounds:
                raise RuntimeError('Spec set was not resolved after %d rounds. '
                        'This is likely a bug.' % max_rounds)

            if not self.resolve_one_round():
                # Break as soon as nothing significant is added in this round
                break

            print 'After round #%d' % round
            print_specset(self.spec_set, round)

        # Return the pinned spec set
        return self.pin_spec_set()

    def pin_spec_set(self):
        """Pins all packages in given resolved spec set and returns a new spec
        set.  Requires the input spec set to be resolved.
        """
        new_spec_set = SpecSet()
        for spec in self.spec_set.normalize():
            best_version = self.pkgmgr.find_best_match(spec)
            pinned_spec = Spec(spec.name, [('==', best_version)], spec.source)
            new_spec_set.add_spec(pinned_spec)
        return new_spec_set

    def find_all_dependencies(self):
        """Finds best matches for the current spec set in the package manager,
        returning all requirements for those packages.
        """
        spec_set = self.spec_set
        pkgmgr = self.pkgmgr

        deps = set()
        for spec in spec_set.normalize():
            version = pkgmgr.find_best_match(spec)
            pkg_deps = pkgmgr.get_dependencies(spec.name, version)

            # Append source information to the new specs
            if spec.source:
                source = '%s ~> %s==%s' % (spec.source, spec.name, version)
            else:
                source = '%s==%s' % (spec.name, version)

            pkg_deps = {s.add_source(source) for s in pkg_deps}
            deps.update(pkg_deps)

        return deps

    def find_new_dependencies(self):
        """Finds all dependencies for the given spec set (in the package
        manager), but only returns what specs are new to the set.
        """
        all_deps = self.find_all_dependencies()
        return all_deps - set(self.spec_set)  # only return _new_ specs
<commit_msg>Make package_manager a required argument.<commit_after>from piptools.datastructures import SpecSet, Spec

def print_specset(specset, round):
    print('After round #%s:' % (round,))
    for spec in specset:
        print('  - %s' % (spec.description(),))

class Resolver(object):
    def __init__(self, spec_set, package_manager):
        """This class resolves a given SpecSet by querying the given
        PackageManager.
        """
        self.spec_set = spec_set
        self.pkgmgr = package_manager

    def resolve_one_round(self):
        """Resolves one level of the current spec set, by finding best matches
        for the current spec set in the package manager and returning all
        (new) requirements for those packages.

        Returns whether the spec set was changed significantly by this round.
        """
        new_deps = self.find_new_dependencies()
        self.spec_set.add_specs(new_deps)
        return len(new_deps) > 0

    def resolve(self, max_rounds=12):
        """Resolves the spec set one round at a time, until the set does not
        change significantly anymore.  Protects against infinite loops by
        breaking out after a max number rounds.
        """
        round = 0
        while True:
            round += 1
            if round > max_rounds:
                raise RuntimeError('Spec set was not resolved after %d rounds. '
                        'This is likely a bug.' % max_rounds)

            if not self.resolve_one_round():
                # Break as soon as nothing significant is added in this round
                break

            print 'After round #%d' % round
            print_specset(self.spec_set, round)

        # Return the pinned spec set
        return self.pin_spec_set()

    def pin_spec_set(self):
        """Pins all packages in given resolved spec set and returns a new spec
        set.  Requires the input spec set to be resolved.
        """
        new_spec_set = SpecSet()
        for spec in self.spec_set.normalize():
            best_version = self.pkgmgr.find_best_match(spec)
            pinned_spec = Spec(spec.name, [('==', best_version)], spec.source)
            new_spec_set.add_spec(pinned_spec)
        return new_spec_set

    def find_all_dependencies(self):
        """Finds best matches for the current spec set in the package manager,
        returning all requirements for those packages.
        """
        spec_set = self.spec_set
        pkgmgr = self.pkgmgr

        deps = set()
        for spec in spec_set.normalize():
            version = pkgmgr.find_best_match(spec)
            pkg_deps = pkgmgr.get_dependencies(spec.name, version)

            # Append source information to the new specs
            if spec.source:
                source = '%s ~> %s==%s' % (spec.source, spec.name, version)
            else:
                source = '%s==%s' % (spec.name, version)

            pkg_deps = {s.add_source(source) for s in pkg_deps}
            deps.update(pkg_deps)

        return deps

    def find_new_dependencies(self):
        """Finds all dependencies for the given spec set (in the package
        manager), but only returns what specs are new to the set.
        """
        all_deps = self.find_all_dependencies()
        return all_deps - set(self.spec_set)  # only return _new_ specs
<|endoftext|>