<commit_before># -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import sys

from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
from airflow.models import BaseOperator
from airflow.models.xcom import MAX_XCOM_SIZE
from airflow.utils.decorators import apply_defaults

class GoogleCloudStorageDownloadOperator(BaseOperator):
    """
    Downloads a file from Google Cloud Storage.

    :param bucket: The Google cloud storage bucket where the object is. (templated)
    :type bucket: str
    :param object: The name of the object to download in the Google cloud
        storage bucket. (templated)
    :type object: str
    :param filename: The file path on the local file system (where the
        operator is being executed) that the file should be downloaded to. (templated)
        If no filename passed, the downloaded data will not be stored on the local file
        system.
    :type filename: str
    :param store_to_xcom_key: If this param is set, the operator will push
        the contents of the downloaded file to XCom with the key set in this
        parameter. If not set, the downloaded data will not be pushed to XCom. (templated)
    :type store_to_xcom_key: str
    :param google_cloud_storage_conn_id: The connection ID to use when
        connecting to Google cloud storage.
    :type google_cloud_storage_conn_id: str
    :param delegate_to: The account to impersonate, if any.
        For this to work, the service account making the request must have
        domain-wide delegation enabled.
    :type delegate_to: str
    """
    template_fields = ('bucket', 'object', 'filename', 'store_to_xcom_key',)
    ui_color = '#f0eee4'

    @apply_defaults
    def __init__(self,
                 bucket,
                 object,
                 filename=None,
                 store_to_xcom_key=None,
                 google_cloud_storage_conn_id='google_cloud_default',
                 delegate_to=None,
                 *args,
                 **kwargs):
        super(GoogleCloudStorageDownloadOperator, self).__init__(*args, **kwargs)
        self.bucket = bucket
        self.object = object
        self.filename = filename
        self.store_to_xcom_key = store_to_xcom_key
        self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
        self.delegate_to = delegate_to

    def execute(self, context):
        self.log.info('Executing download: %s, %s, %s', self.bucket,
                      self.object, self.filename)
        hook = GoogleCloudStorageHook(
            google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
            delegate_to=self.delegate_to
        )
        file_bytes = hook.download(bucket=self.bucket,
                                   object=self.object,
                                   filename=self.filename)
        if self.store_to_xcom_key:
            if sys.getsizeof(file_bytes) < MAX_XCOM_SIZE:
                context['ti'].xcom_push(key=self.store_to_xcom_key, value=file_bytes)
            else:
                raise RuntimeError(
                    'The size of the downloaded file is too large to push to XCom!'
                )
<commit_msg>Add clarity to gcs_download_operator params<commit_after># -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import sys

from airflow.contrib.hooks.gcs_hook import GoogleCloudStorageHook
from airflow.models import BaseOperator
from airflow.models.xcom import MAX_XCOM_SIZE
from airflow.utils.decorators import apply_defaults

class GoogleCloudStorageDownloadOperator(BaseOperator):
    """
    Downloads a file from Google Cloud Storage.

    If a filename is supplied, it writes the file to the specified location, alternatively one can
    set the ``store_to_xcom_key`` parameter to True push the file content into xcom. When the file size
    exceeds the maximum size for xcom it is recommended to write to a file.

    :param bucket: The Google cloud storage bucket where the object is.
        Must not contain 'gs://' prefix. (templated)
    :type bucket: str
    :param object: The name of the object to download in the Google cloud
        storage bucket. (templated)
    :type object: str
    :param filename: The file path, including filename,  on the local file system (where the
        operator is being executed) that the file should be downloaded to. (templated)
        If no filename passed, the downloaded data will not be stored on the local file
        system.
    :type filename: str
    :param store_to_xcom_key: If this param is set, the operator will push
        the contents of the downloaded file to XCom with the key set in this
        parameter. If not set, the downloaded data will not be pushed to XCom. (templated)
    :type store_to_xcom_key: str
    :param google_cloud_storage_conn_id: The connection ID to use when
        connecting to Google cloud storage.
    :type google_cloud_storage_conn_id: str
    :param delegate_to: The account to impersonate, if any.
        For this to work, the service account making the request must have
        domain-wide delegation enabled.
    :type delegate_to: str
    """
    template_fields = ('bucket', 'object', 'filename', 'store_to_xcom_key',)
    ui_color = '#f0eee4'

    @apply_defaults
    def __init__(self,
                 bucket,
                 object,
                 filename=None,
                 store_to_xcom_key=None,
                 google_cloud_storage_conn_id='google_cloud_default',
                 delegate_to=None,
                 *args,
                 **kwargs):
        super(GoogleCloudStorageDownloadOperator, self).__init__(*args, **kwargs)
        self.bucket = bucket
        self.object = object
        self.filename = filename
        self.store_to_xcom_key = store_to_xcom_key
        self.google_cloud_storage_conn_id = google_cloud_storage_conn_id
        self.delegate_to = delegate_to

    def execute(self, context):
        self.log.info('Executing download: %s, %s, %s', self.bucket,
                      self.object, self.filename)
        hook = GoogleCloudStorageHook(
            google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
            delegate_to=self.delegate_to
        )
        file_bytes = hook.download(bucket=self.bucket,
                                   object=self.object,
                                   filename=self.filename)
        if self.store_to_xcom_key:
            if sys.getsizeof(file_bytes) < MAX_XCOM_SIZE:
                context['ti'].xcom_push(key=self.store_to_xcom_key, value=file_bytes)
            else:
                raise RuntimeError(
                    'The size of the downloaded file is too large to push to XCom!'
                )
<|endoftext|>


<commit_before>def curved_background_subtraction_2d(data2d, mask2d):

    import numpy
    n_col = numpy.size(data2d[0:1, :])
    n_row = numpy.size(data2d[:, 0:1])
    avg_bkgr_tot = 0.0
    avg_bkgr_cont = 0
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if mask2d[row, col] == 1:

                loc_bkgr_tot = 0.0
                loc_bkgr_cont = 0
                if mask2d[n_row - 1, col] == 0:
                    loc_bkgr_tot += data2d[n_row - 1, col]
                    loc_bkgr_cont += 1
                if mask2d[0, col] == 0:
                    loc_bkgr_tot += data2d[0, col]
                    loc_bkgr_cont += 1
                if mask2d[row, n_col - 1] == 0:
                    loc_bkgr_tot += data2d[row, n_col - 1]
                    loc_bkgr_cont += 1
                if mask2d[row, 0] == 0:
                    loc_bkgr_tot += data2d[row, 0]
                    loc_bkgr_cont += 1
                loc_bkgr = loc_bkgr_tot / float(loc_bkgr_cont)

                avg_bkgr_tot += loc_bkgr
                avg_bkgr_cont += 1

                if data2d[row, col] > loc_bkgr:
                    data2d[row, col] = data2d[row, col] - loc_bkgr
                else:
                    data2d[row, col] = 0
    avg_bkgr = avg_bkgr_tot / float(avg_bkgr_cont)
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if mask2d[row, col] == 0:
                data2d[row, col] = 0
    return avg_bkgr

def flat_background_subtraction_2d(data2d, diffdata2d_ext):

    import numpy
    n_col = numpy.size(data2d[0:1, :])
    n_row = numpy.size(data2d[:, 0:1])

    tot_bkgr = 0.0
    cont = 0.0
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if diffdata2d_ext[row, col] == 0:
                cont += 1
                tot_bkgr += data2d[row, col]
    if tot_bkgr > 0 and cont > 0:
        bkgr = tot_bkgr / cont
    else:
        bkgr = 0
    #print 'bkgr=', bkgr
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if diffdata2d_ext[row, col] == 1 and data2d[row, col] > bkgr:
                data2d[row, col] = data2d[row, col] - bkgr
            else:
                data2d[row, col] = 0
    return bkgr
<commit_msg>fixing the division by zero issue<commit_after>def curved_background_subtraction_2d(data2d, mask2d):

    import numpy
    n_col = numpy.size(data2d[0:1, :])
    n_row = numpy.size(data2d[:, 0:1])
    avg_bkgr_tot = 0.0
    avg_bkgr_cont = 0
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if mask2d[row, col] == 1:

                loc_bkgr_tot = 0.0
                loc_bkgr_cont = 0
                if mask2d[n_row - 1, col] == 0:
                    loc_bkgr_tot += data2d[n_row - 1, col]
                    loc_bkgr_cont += 1
                if mask2d[0, col] == 0:
                    loc_bkgr_tot += data2d[0, col]
                    loc_bkgr_cont += 1
                if mask2d[row, n_col - 1] == 0:
                    loc_bkgr_tot += data2d[row, n_col - 1]
                    loc_bkgr_cont += 1
                if mask2d[row, 0] == 0:
                    loc_bkgr_tot += data2d[row, 0]
                    loc_bkgr_cont += 1

                if loc_bkgr_cont > 0:
                    loc_bkgr = loc_bkgr_tot / float(loc_bkgr_cont)
                else:
                    loc_bkgr = 0
                avg_bkgr_tot += loc_bkgr
                avg_bkgr_cont += 1

                if data2d[row, col] > loc_bkgr:
                    data2d[row, col] = data2d[row, col] - loc_bkgr
                else:
                    data2d[row, col] = 0
    avg_bkgr = avg_bkgr_tot / float(avg_bkgr_cont)
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if mask2d[row, col] == 0:
                data2d[row, col] = 0
    return avg_bkgr

def flat_background_subtraction_2d(data2d, diffdata2d_ext):

    import numpy
    n_col = numpy.size(data2d[0:1, :])
    n_row = numpy.size(data2d[:, 0:1])

    tot_bkgr = 0.0
    cont = 0.0
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if diffdata2d_ext[row, col] == 0:
                cont += 1
                tot_bkgr += data2d[row, col]
    if tot_bkgr > 0 and cont > 0:
        bkgr = tot_bkgr / cont
    else:
        bkgr = 0
    #print 'bkgr=', bkgr
    for row in range(0, n_row, 1):
        for col in range(0, n_col, 1):
            if diffdata2d_ext[row, col] == 1 and data2d[row, col] > bkgr:
                data2d[row, col] = data2d[row, col] - bkgr
            else:
                data2d[row, col] = 0
    return bkgr
<|endoftext|>


<commit_before>from django.core.management.base import BaseCommand, AppCommand
from django.conf import settings
from optparse import make_option
import os
import glob
from lingcod.layers.models import PrivateKml
from django.contrib.auth.models import User, Group

class Command(BaseCommand):
    help = "Populates the PrivateKml table from the PRIVATE_KML_ROOT contents .. a good starting point"
    args = '[optional group name to share all KMLs with]'

    # Validation is called explicitly each time the server is reloaded.
    requires_model_validation = False

    def handle(self, groupname=None, *args, **options):
        for pkml in PrivateKml.objects.all():
            pkml.delete()

        if groupname:
            g = Group.objects.get(name=groupname)

        if not os.path.exists(settings.PRIVATE_KML_ROOT):
            raise Exception("Please create or set up a PRIVATE_KML_ROOT directory (currently set to %s" %
                    settings.PRIVATE_KML_ROOT)
        for d in os.listdir(settings.PRIVATE_KML_ROOT):
            path = os.path.join(settings.PRIVATE_KML_ROOT,d)
            kmls = glob.glob(os.path.join(path,'*.km*'))
            if len(kmls) == 0:
                print "No KML/KMZ found in %s" % path
                continue

            for kml in kmls:
                basename = os.path.basename(kml).split('.')[0]
                privatekml_name = d+'_'+basename
                try:
                    pkml = PrivateKml.objects.create(name=privatekml_name[:99],base_kml=kml)
                    if groupname:
                        pkml.sharing_groups.add(g)
                    print "Created %s from %s" % (pkml,kml)
                except Exception as e:
                    print "couldn't create privatekml from %s" % kml
                    print "  ", e
<commit_msg>Limit to km[z,l]$ files<commit_after>from django.core.management.base import BaseCommand, AppCommand
from django.conf import settings
from optparse import make_option
import os
import glob
from lingcod.layers.models import PrivateKml
from django.contrib.auth.models import User, Group

class Command(BaseCommand):
    help = "Populates the PrivateKml table from the PRIVATE_KML_ROOT contents .. a good starting point"
    args = '[optional group name to share all KMLs with]'

    # Validation is called explicitly each time the server is reloaded.
    requires_model_validation = False

    def handle(self, groupname=None, *args, **options):
        for pkml in PrivateKml.objects.all():
            pkml.delete()

        if groupname:
            g = Group.objects.get(name=groupname)

        if not os.path.exists(settings.PRIVATE_KML_ROOT):
            raise Exception("Please create or set up a PRIVATE_KML_ROOT directory (currently set to %s" %
                    settings.PRIVATE_KML_ROOT)
        for d in os.listdir(settings.PRIVATE_KML_ROOT):
            path = os.path.join(settings.PRIVATE_KML_ROOT,d)
            kmls = glob.glob(os.path.join(path,'*.km[z,l]'))
            if len(kmls) == 0:
                print "No KML/KMZ found in %s" % path
                continue

            for kml in kmls:
                basename = os.path.basename(kml).split('.')[0]
                privatekml_name = d+'_'+basename
                print "Creating", privatekml_name
                try:
                    pkml = PrivateKml.objects.create(name=privatekml_name[:99],base_kml=kml)
                    if groupname:
                        pkml.sharing_groups.add(g)
                    print " created %s from %s" % (pkml,kml)
                except Exception as e:
                    print " couldn't create privatekml from %s" % kml
                    print "  ", e
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Tue May  5 21:15:28 2015

@author: richard

creates a trajectory object that has a bunch of sweet methods.
takes vectors from agent and adds it to ensemble

trajectory.ensemble
trajectory.append_ensemble(a trajectory)

trajectory.describe()
    agent info
    plots

trajectory.save
"""
import pandas as pd
import plotting_funcs3D
import numpy as np

def T_find_stats(t_targfinds):
    """Target finding stats. Returns average time to find the target, and the
    number of successes

    TODO:  find tfind stats for ensemble
    """
    t_targfinds = np.array(t_targfinds)
    t_finds_NoNaNs = t_targfinds[~np.isnan(t_targfinds)]  # remove NaNs
    if len(t_finds_NoNaNs) == 0:

        return np.nan  # no target finds, no stats to run
    else:
        num_success = float(len(t_finds_NoNaNs))
        Tfind_avg = sum(t_finds_NoNaNs)/len(t_finds_NoNaNs)

        return Tfind_avg

class Trajectory():
    def __init__(self):
        self.ensemble = pd.DataFrame()
        self.agent_info = {}

    def append_ensemble(self, arraydict):
        trajectory = pd.DataFrame(arraydict)
        self.ensemble = self.ensemble.append(trajectory)

    def plot_single_trajectory(self, trajectory_i=0):
        plot_kwargs = {'title':"Individual agent trajectory", 'titleappend':' (id = {})'.format(trajectory_i)}
        plotting_funcs3D.plot_single_trajectory(self.ensemble.loc[self.ensemble['trajectory_num']==trajectory_i], self.agent_info, plot_kwargs)

    def add_agent_info(self, data_dict):
        for key, value in data_dict.iteritems():
            self.agent_info[key] = value

    def plot_force_violin(self):
        plotting_funcs3D.force_violin(self.ensemble, self.agent_info)

    def plot_posheatmap(self):
        plotting_funcs3D.heatmaps(self.ensemble, self.agent_info)

    def plot_kinematic_hists(self, ensemble='none', titleappend=''):
        print titleappend
        if type(ensemble) is str:
            ensemble = self.ensemble
            ensemble = ensemble.loc[(ensemble['position_x'] >0.25) & (ensemble['position_x'] <0.95)]
        plotting_funcs3D.stateHistograms(ensemble, self.agent_info, titleappend)

    def plot_door_velocity_compass(self, region='door', kind='bin_average'):

        if region == 'door':
            """plot the area """
            ensemble = self.ensemble.loc[((self.ensemble['position_x']>0.25) & (self.ensemble['position_x']<0.5)), ['velocity_x', 'velocity_y']]
        else:
            ensemble = self.ensemble

        plotting_funcs3D.compass_plots(ensemble, self.agent_info, kind)

    def plot_single_3Dtrajectory(self, trajectory_i=0):
        plot_kwargs = {'title':"Individual agent trajectory", 'titleappend':' (id = {})'.format(trajectory_i)}
        plotting_funcs3D.plot3D_trajectory(self.ensemble.loc[self.ensemble['trajectory_num']==trajectory_i], self.agent_info, plot_kwargs)

    def plot_sliced_hists(self):
        """Plot histograms from 0.25 < x < 0.95, as well as that same space
        divided into 4 equal 0.15 m segments
        """
        x_edges = np.linspace(.25,.95,5)

        print "full ensemble"
        full_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
            & (self.ensemble['position_x'] < x_edges[4])]
        self.plot_kinematic_hists(ensemble=full_ensemble, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[4]))

        print "downwind half ensemble"
        downwind_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
            & (self.ensemble['position_x'] < x_edges[2])]
        self.plot_kinematic_hists(ensemble=downwind_ensemble, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[2]))

        print "upwind half ensemble"
        upwind_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[2]) \
            & (self.ensemble['position_x'] < x_edges[4])]
        self.plot_kinematic_hists(ensemble=upwind_ensemble, titleappend=' {} < x <= {}'.format(x_edges[2], x_edges[4]))
#
#        print "first quarter ensemble"
#        ensemble1 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
#            & (self.ensemble['position_x'] <= x_edges[1])]
#        self.plot_kinematic_hists(ensemble=ensemble1, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[1]))
#
#        print "second quarter ensemble"
#        ensemble2 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[1]) \
#            & (self.ensemble['position_x'] <= x_edges[2])]
#        self.plot_kinematic_hists(ensemble=ensemble2, titleappend=' {} < x <= {}'.format(x_edges[1], x_edges[2]))
#
#        print "third quarter ensemble"
#        ensemble3 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[2]) \
#            & (self.ensemble['position_x'] <= x_edges[3])]
#        self.plot_kinematic_hists(ensemble=ensemble3, titleappend=' {} < x <= {}'.format(x_edges[2], x_edges[3]))
#
#        print "fourth quarter ensemble"
#        ensemble4 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[3]) \
#            & (self.ensemble['position_x'] <= x_edges[4])]
#        self.plot_kinematic_hists(ensemble=ensemble4, titleappend=' {} < x <= {}'.format(x_edges[3], x_edges[4]))
#
        print "full- + up- + down-wind"
        plotting_funcs3D.stateHistograms(full_ensemble, self.agent_info, titleappend = '', upw_ensemble = upwind_ensemble, downw_ensemble = downwind_ensemble)

    def describe(self, plot_kwargs={'trajectories':False, 'heatmap':True, 'states':True, 'singletrajectories':False, 'force_violin':True}):
        print self.ensemble.describe()
        if plot_kwargs['heatmap'] is True:
            self.plot_posheatmap()
        if plot_kwargs['singletrajectories'] is True:
            self.plot_single_trajectory()
        if plot_kwargs['states'] is True:
    #         plot histogram of pos, velo, accel distributions
            self.plot_kinematic_hists()
        if plot_kwargs['force_violin'] is True:
            self.plot_force_violin()

    def dump2csvs(self):
        for trajectory_i in self.ensemble.trajectory_num.unique():
            temp_traj = self.ensemble[self.ensemble['trajectory_num'] == trajectory_i]
            temp_array = temp_traj[['position_x', 'position_y']].values
            np.savetxt(str(trajectory_i) + ".csv", temp_array, delimiter=",")

#    def <commit_msg>solving plume in/out stats, adding compass plot of all forces func<commit_after># -*- coding: utf-8 -*-
"""
Created on Tue May  5 21:15:28 2015

@author: richard

creates a trajectory object that has a bunch of sweet methods.
takes vectors from agent and adds it to ensemble

trajectory.ensemble
trajectory.append_ensemble(a trajectory)

trajectory.describe()
    agent info
    plots

trajectory.save
"""
import pandas as pd
import plotting_funcs3D
import numpy as np

def T_find_stats(t_targfinds):
    """Target finding stats. Returns average time to find the target, and the
    number of successes

    TODO:  find tfind stats for ensemble
    """
    t_targfinds = np.array(t_targfinds)
    t_finds_NoNaNs = t_targfinds[~np.isnan(t_targfinds)]  # remove NaNs
    if len(t_finds_NoNaNs) == 0:

        return np.nan  # no target finds, no stats to run
    else:
        num_success = float(len(t_finds_NoNaNs))
        Tfind_avg = sum(t_finds_NoNaNs)/len(t_finds_NoNaNs)

        return Tfind_avg

class Trajectory():
    def __init__(self):
        self.ensemble = pd.DataFrame()
        self.agent_info = {}

    def append_ensemble(self, arraydict):
        trajectory = pd.DataFrame(arraydict)
        self.ensemble = self.ensemble.append(trajectory)

    def plot_single_trajectory(self, trajectory_i=0):
        plot_kwargs = {'title':"Individual agent trajectory", 'titleappend':' (id = {})'.format(trajectory_i)}
        plotting_funcs3D.plot_single_trajectory(self.ensemble.loc[self.ensemble['trajectory_num']==trajectory_i], self.agent_info, plot_kwargs)

    def add_agent_info(self, metadata_dict):
        for key, value in metadata_dict.iteritems():
            self.agent_info[key] = value

    def plot_force_violin(self):
        plotting_funcs3D.force_violin(self.ensemble, self.agent_info)

    def plot_posheatmap(self):
        plotting_funcs3D.heatmaps(self.ensemble, self.agent_info)

    def plot_kinematic_hists(self, ensemble='none', titleappend=''):
        print titleappend
        if type(ensemble) is str:
            ensemble = self.ensemble
            ensemble = ensemble.loc[(ensemble['position_x'] >0.25) & (ensemble['position_x'] <0.95)]
        plotting_funcs3D.stateHistograms(ensemble, self.agent_info, titleappend)

    def plot_door_velocity_compass(self, region='door', kind='bin_average'):

        if region == 'door':
            """plot the area """
            ensemble = self.ensemble.loc[((self.ensemble['position_x']>0.25) & (self.ensemble['position_x']<0.5)), ['velocity_x', 'velocity_y']]
        else:
            ensemble = self.ensemble

        plotting_funcs3D.compass_histogram(ensemble, self.agent_info, kind)

    def plot_compass(self, kind='bin_average'):
        for force in self.agent_info['forces']:
            magnitudes, thetas = getattr(self.ensemble, name+'_xy_mag').values, getattr(V, name+'_xy_theta').values

#        plotting_funcs3D.compass_histogram(ensemble, self.agent_info, kind)

    def plot_single_3Dtrajectory(self, trajectory_i=0):
        plot_kwargs = {'title':"Individual agent trajectory", 'titleappend':' (id = {})'.format(trajectory_i)}
        plotting_funcs3D.plot3D_trajectory(self.ensemble.loc[self.ensemble['trajectory_num']==trajectory_i], self.agent_info, plot_kwargs)

    def plot_sliced_hists(self):
        """Plot histograms from 0.25 < x < 0.95, as well as that same space
        divided into 4 equal 0.15 m segments
        """
        x_edges = np.linspace(.25,.95,5)

        print "full ensemble"
        full_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
            & (self.ensemble['position_x'] < x_edges[4])]
        self.plot_kinematic_hists(ensemble=full_ensemble, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[4]))

        print "downwind half ensemble"
        downwind_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
            & (self.ensemble['position_x'] < x_edges[2])]
        self.plot_kinematic_hists(ensemble=downwind_ensemble, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[2]))

        print "upwind half ensemble"
        upwind_ensemble = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[2]) \
            & (self.ensemble['position_x'] < x_edges[4])]
        self.plot_kinematic_hists(ensemble=upwind_ensemble, titleappend=' {} < x <= {}'.format(x_edges[2], x_edges[4]))
#
#        print "first quarter ensemble"
#        ensemble1 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[0]) \
#            & (self.ensemble['position_x'] <= x_edges[1])]
#        self.plot_kinematic_hists(ensemble=ensemble1, titleappend=' {} < x <= {}'.format(x_edges[0], x_edges[1]))
#
#        print "second quarter ensemble"
#        ensemble2 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[1]) \
#            & (self.ensemble['position_x'] <= x_edges[2])]
#        self.plot_kinematic_hists(ensemble=ensemble2, titleappend=' {} < x <= {}'.format(x_edges[1], x_edges[2]))
#
#        print "third quarter ensemble"
#        ensemble3 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[2]) \
#            & (self.ensemble['position_x'] <= x_edges[3])]
#        self.plot_kinematic_hists(ensemble=ensemble3, titleappend=' {} < x <= {}'.format(x_edges[2], x_edges[3]))
#
#        print "fourth quarter ensemble"
#        ensemble4 = self.ensemble.loc[(self.ensemble['position_x'] > x_edges[3]) \
#            & (self.ensemble['position_x'] <= x_edges[4])]
#        self.plot_kinematic_hists(ensemble=ensemble4, titleappend=' {} < x <= {}'.format(x_edges[3], x_edges[4]))
#
        print "full- + up- + down-wind"
        plotting_funcs3D.stateHistograms(full_ensemble, self.agent_info, titleappend = '', upw_ensemble = upwind_ensemble, downw_ensemble = downwind_ensemble)

    def plume_stats(self):
        """
        in plume == 1, out == 0. therefore sum/n is % in plume
        """
        in_out = self.ensemble.inPlume.values
        print "Total timesteps: {size}. Total time in plume: {vecsum}. Ratio:"\
            "{ratio}".format(size=in_out.size, vecsum=in_out.sum(), ratio=in_out.sum()/in_out.size)

    def describe(self, plot_kwargs={'trajectories':False, 'heatmap':True, 'states':True, 'singletrajectories':False, 'force_violin':True}):
        print self.ensemble.describe()
        if plot_kwargs['heatmap'] is True:
            self.plot_posheatmap()
        if plot_kwargs['singletrajectories'] is True:
            self.plot_single_trajectory()
        if plot_kwargs['states'] is True:
    #         plot histogram of pos, velo, accel distributions
            self.plot_kinematic_hists()
        if plot_kwargs['force_violin'] is True:
            self.plot_force_violin()

    def dump2csvs(self):
        for trajectory_i in self.ensemble.trajectory_num.unique():
            temp_traj = self.ensemble[self.ensemble['trajectory_num'] == trajectory_i]
            temp_array = temp_traj[['position_x', 'position_y']].values
            np.savetxt(str(trajectory_i) + ".csv", temp_array, delimiter=",")

#    def <|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import unittest

from . import address_request
from ..examples import address
from ..types import Object

class TestPrimitiveType(unittest.TestCase):

    def test_object(self):
        obj = Object(properties=address.Address)
        inst = obj.bind(address_request)

        self.assertIsInstance(inst, address.Address)
        self.assertEqual(inst.country_name, 'Japan')
        self.assertEqual(inst.region, 'Tokyo')
        self.assertEqual(inst.locality, 'Chiyoda-ku')
        self.assertEqual(inst.street_address, '1-3-6 Kudan-kita')
        self.assertEqual(inst.extended_address, '7F')

        self.assertIsNone(inst.postal_code)
        self.assertIsNone(inst.post_office_box)

    def test_object_without_properties(self):
        obj = Object()
        inst = obj.bind(address_request)

        self.assertIs(inst, address_request)
<commit_msg>fix testcase<commit_after># -*- coding: utf-8 -*-

import unittest

from . import product_request
from ..examples import product
from ..types import Object

class TestPrimitiveType(unittest.TestCase):

    def test_object(self):
        obj = Object(properties=product.Product)
        inst = obj.bind(product_request)

        self.assertIsInstance(inst, product.Product)
        self.assertEqual(inst.id, 2)
        self.assertEqual(inst.name, "An ice sculpture")
        self.assertEqual(inst.price, 12.50)
        self.assertEqual(inst.tags, ["cold", "ice"])
        self.assertEqual(inst.dimensions.length, 7.0)
        self.assertEqual(inst.dimensions.width, 12.0)
        self.assertEqual(inst.dimensions.height, 9.5)
        self.assertEqual(inst.warehouseLocation, {
            "latitude": -78.75,
            "longitude": 20.4,
        })
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import unicode_literals

__version__ = '1.0.12'
<commit_msg>release 1.1<commit_after># -*- coding: utf-8 -*-

from __future__ import unicode_literals

__version__ = '1.1'
<|endoftext|>


<commit_before>###

"""Functionality common to all account classes.

Each subclass of Account must provide five methods:
  (*) create() and destroy(), which are static;
  (*) configure(), start(), and stop(), which are not.

configure(), which takes a record as its only argument, does
things like set up ssh keys. In addition, an Account subclass must
provide static member variables SHELL, which contains the unique shell
that it uses; and TYPE, a string that is used by the account creation
code.  For no particular reason, TYPE is divided hierarchically by
periods; at the moment the only convention is that all sliver accounts
have type that begins with sliver.

There are any number of race conditions that may result from the fact
that account names are not unique over time.  Moreover, it's a bad
idea to perform lengthy operations while holding the database lock.
In order to deal with both of these problems, we use a worker thread
for each account name that ever exists.  On 32-bit systems with large
numbers of accounts, this may cause the NM process to run out of
*virtual* memory!  This problem may be remedied by decreasing the
maximum stack size.
"""

import os
import pwd, grp
import threading

import logger
import tools

# shell path -> account class association
shell_acct_class = {}
# account type -> account class association
type_acct_class = {}

# these semaphores are acquired before creating/destroying an account
create_sem = threading.Semaphore(1)
destroy_sem = threading.Semaphore(1)

def register_class(acct_class):
    """Call once for each account class. This method adds the class
to the dictionaries used to look up account classes by shell and
type."""
    shell_acct_class[acct_class.SHELL] = acct_class
    type_acct_class[acct_class.TYPE] = acct_class

# private account name -> worker object association and associated lock
name_worker_lock = threading.Lock()
name_worker = {}

def allpwents():
    return [pw_ent for pw_ent in pwd.getpwall() if pw_ent[6] in shell_acct_class]

def all():
    """Return the names of all accounts on the system with recognized shells."""
    return [pw_ent[0] for pw_ent in allpwents()]

def get(name):
    """Return the worker object for a particular username.  If no such object exists, create it first."""
    name_worker_lock.acquire()
    try:
        if name not in name_worker: name_worker[name] = Worker(name)
        return name_worker[name]
    finally: name_worker_lock.release()

# xxx strictly speaking this class should not use self.name that in fact
# is accidentally inited by the subclasses constructor...
class Account:
    def __init__(self, name):
        self.name = name
        self.keys = ''
        logger.verbose('account: Initing account %s'%name)

    @staticmethod
    def create(name, vref = None): abstract

    @staticmethod
    def destroy(name): abstract

    def configure(self, rec):
        """Write <rec['keys']> to my authorized_keys file."""
        logger.verbose('account: configuring %s'%self.name)
        new_keys = rec['keys']
        if new_keys != self.keys:
            # get the unix account info
            gid = grp.getgrnam("slices")[2]
            pw_info = pwd.getpwnam(self.name)
            uid = pw_info[2]
            pw_dir = pw_info[5]

            # write out authorized_keys file and conditionally create
            # the .ssh subdir if need be.
            dot_ssh = os.path.join(pw_dir,'.ssh')
            if not os.path.isdir(dot_ssh):
                if not os.path.isdir(pw_dir):
                    logger.verbose('account: WARNING: homedir %s does not exist for %s!'%(pw_dir,self.name))
                    os.mkdir(pw_dir)
                    os.chown(pw_dir, uid, gid)
                os.mkdir(dot_ssh)

            auth_keys = os.path.join(dot_ssh,'authorized_keys')
            tools.write_file(auth_keys, lambda f: f.write(new_keys))

            # set access permissions and ownership properly
            os.chmod(dot_ssh, 0700)
            os.chown(dot_ssh, uid, gid)
            os.chmod(auth_keys, 0600)
            os.chown(auth_keys, uid, gid)

            # set self.keys to new_keys only when all of the above ops succeed
            self.keys = new_keys

            logger.log('account: %s: installed ssh keys' % self.name)

    def start(self, delay=0): pass
    def stop(self): pass
    def is_running(self): pass

class Worker:

    def __init__(self, name):
        self.name = name  # username
        self._acct = None  # the account object currently associated with this worker

    def ensure_created(self, rec):
        """Check account type is still valid.  If not, recreate sliver.
If still valid, check if running and configure/start if not."""
        logger.log_data_in_file(rec,"/var/lib/nodemanager/%s.rec.txt"%rec['name'],
                                'raw rec captured in ensure_created',logger.LOG_VERBOSE)
        curr_class = self._get_class()
        next_class = type_acct_class[rec['type']]
        if next_class != curr_class:
            self._destroy(curr_class)
            create_sem.acquire()
            try: next_class.create(self.name, rec)
            finally: create_sem.release()
        if not isinstance(self._acct, next_class): self._acct = next_class(rec)
        logger.verbose("account.Worker.ensure_created: %s, running=%r"%(self.name,self.is_running()))

        # reservation_alive is set on reervable nodes, and its value is a boolean
        if 'reservation_alive' in rec:
            # reservable nodes
            if rec['reservation_alive']:
                # this sliver has the lease, it is safe to start it
                if not self.is_running(): self.start(rec)
                else: self.configure(rec)
            else:
                # not having the lease, do not start it
                self.configure(rec)
        # usual nodes - preserve old code
        # xxx it's not clear what to do when a sliver changes type/class
        # in a reservable node
        else:
            if not self.is_running() or next_class != curr_class:
                self.start(rec)
            else: self.configure(rec)

    def ensure_destroyed(self): self._destroy(self._get_class())

    def start(self, rec, d = 0):
        self._acct.configure(rec)
        self._acct.start(delay=d)

    def configure(self, rec):
        self._acct.configure(rec)

    def stop(self): self._acct.stop()

    def is_running(self):
        if (self._acct != None) and self._acct.is_running():
            status = True
        else:
            status = False
            logger.verbose("account: Worker(%s): is not running" % self.name)
        return status

    def _destroy(self, curr_class):
        self._acct = None
        if curr_class:
            destroy_sem.acquire()
            try: curr_class.destroy(self.name)
            finally: destroy_sem.release()

    def _get_class(self):
        try: shell = pwd.getpwnam(self.name)[6]
        except KeyError: return None
        return shell_acct_class[shell]

    # bind mount root side dir to sliver side
    # needs to be done before sliver starts, in the vserver case at least
    def expose_ssh_dir (self):
        try:
            root_ssh="/home/%s/.ssh"%self.name
            sliver_ssh="/vservers/%s/home/%s/.ssh"%(self.name,self.name)
            # any of both might not exist yet
            for path in [root_ssh,sliver_ssh]:
                if not os.path.exists (path):
                    os.mkdir(path)
                if not os.path.isdir (path):
                    raise Exception
            mounts=file('/proc/mounts').read()
            if mounts.find(sliver_ssh)<0:
                # xxx perform mount
                subprocess.call("mount --bind -o ro %s %s"%(root_ssh,sliver_ssh),shell=True)
                logger.log("expose_ssh_dir: %s mounted into slice %s"%(root_ssh,self.name))
        except:
            logger.log_exc("expose_ssh_dir with slice %s failed"%self.name)

<commit_msg>expose_ssh_dir relocated from Worker to the Account class where it belongs<commit_after>###

"""Functionality common to all account classes.

Each subclass of Account must provide five methods:
  (*) create() and destroy(), which are static;
  (*) configure(), start(), and stop(), which are not.

configure(), which takes a record as its only argument, does
things like set up ssh keys. In addition, an Account subclass must
provide static member variables SHELL, which contains the unique shell
that it uses; and TYPE, a string that is used by the account creation
code.  For no particular reason, TYPE is divided hierarchically by
periods; at the moment the only convention is that all sliver accounts
have type that begins with sliver.

There are any number of race conditions that may result from the fact
that account names are not unique over time.  Moreover, it's a bad
idea to perform lengthy operations while holding the database lock.
In order to deal with both of these problems, we use a worker thread
for each account name that ever exists.  On 32-bit systems with large
numbers of accounts, this may cause the NM process to run out of
*virtual* memory!  This problem may be remedied by decreasing the
maximum stack size.
"""

import os
import pwd, grp
import threading

import logger
import tools

# shell path -> account class association
shell_acct_class = {}
# account type -> account class association
type_acct_class = {}

# these semaphores are acquired before creating/destroying an account
create_sem = threading.Semaphore(1)
destroy_sem = threading.Semaphore(1)

def register_class(acct_class):
    """Call once for each account class. This method adds the class
to the dictionaries used to look up account classes by shell and
type."""
    shell_acct_class[acct_class.SHELL] = acct_class
    type_acct_class[acct_class.TYPE] = acct_class

# private account name -> worker object association and associated lock
name_worker_lock = threading.Lock()
name_worker = {}

def allpwents():
    return [pw_ent for pw_ent in pwd.getpwall() if pw_ent[6] in shell_acct_class]

def all():
    """Return the names of all accounts on the system with recognized shells."""
    return [pw_ent[0] for pw_ent in allpwents()]

def get(name):
    """Return the worker object for a particular username.  If no such object exists, create it first."""
    name_worker_lock.acquire()
    try:
        if name not in name_worker: name_worker[name] = Worker(name)
        return name_worker[name]
    finally: name_worker_lock.release()

# xxx strictly speaking this class should not use self.name that in fact
# is accidentally inited by the subclasses constructor...
class Account:
    def __init__(self, name):
        self.name = name
        self.keys = ''
        logger.verbose('account: Initing account %s'%name)

    @staticmethod
    def create(name, vref = None): abstract

    @staticmethod
    def destroy(name): abstract

    def configure(self, rec):
        """Write <rec['keys']> to my authorized_keys file."""
        logger.verbose('account: configuring %s'%self.name)
        new_keys = rec['keys']
        if new_keys != self.keys:
            # get the unix account info
            gid = grp.getgrnam("slices")[2]
            pw_info = pwd.getpwnam(self.name)
            uid = pw_info[2]
            pw_dir = pw_info[5]

            # write out authorized_keys file and conditionally create
            # the .ssh subdir if need be.
            dot_ssh = os.path.join(pw_dir,'.ssh')
            if not os.path.isdir(dot_ssh):
                if not os.path.isdir(pw_dir):
                    logger.verbose('account: WARNING: homedir %s does not exist for %s!'%(pw_dir,self.name))
                    os.mkdir(pw_dir)
                    os.chown(pw_dir, uid, gid)
                os.mkdir(dot_ssh)

            auth_keys = os.path.join(dot_ssh,'authorized_keys')
            tools.write_file(auth_keys, lambda f: f.write(new_keys))

            # set access permissions and ownership properly
            os.chmod(dot_ssh, 0700)
            os.chown(dot_ssh, uid, gid)
            os.chmod(auth_keys, 0600)
            os.chown(auth_keys, uid, gid)

            # set self.keys to new_keys only when all of the above ops succeed
            self.keys = new_keys

            logger.log('account: %s: installed ssh keys' % self.name)

    def start(self, delay=0): pass
    def stop(self): pass
    def is_running(self): pass

    # bind mount root side dir to sliver side
    # needs to be done before sliver starts, in the vserver case at least
    def expose_ssh_dir (self):
        try:
            root_ssh="/home/%s/.ssh"%self.name
            sliver_ssh="/vservers/%s/home/%s/.ssh"%(self.name,self.name)
            # any of both might not exist yet
            for path in [root_ssh,sliver_ssh]:
                if not os.path.exists (path):
                    os.mkdir(path)
                if not os.path.isdir (path):
                    raise Exception
            mounts=file('/proc/mounts').read()
            if mounts.find(sliver_ssh)<0:
                # xxx perform mount
                subprocess.call("mount --bind -o ro %s %s"%(root_ssh,sliver_ssh),shell=True)
                logger.log("expose_ssh_dir: %s mounted into slice %s"%(root_ssh,self.name))
        except:
            logger.log_exc("expose_ssh_dir with slice %s failed"%self.name)

class Worker:

    def __init__(self, name):
        self.name = name  # username
        self._acct = None  # the account object currently associated with this worker

    def ensure_created(self, rec):
        """Check account type is still valid.  If not, recreate sliver.
If still valid, check if running and configure/start if not."""
        logger.log_data_in_file(rec,"/var/lib/nodemanager/%s.rec.txt"%rec['name'],
                                'raw rec captured in ensure_created',logger.LOG_VERBOSE)
        curr_class = self._get_class()
        next_class = type_acct_class[rec['type']]
        if next_class != curr_class:
            self._destroy(curr_class)
            create_sem.acquire()
            try: next_class.create(self.name, rec)
            finally: create_sem.release()
        if not isinstance(self._acct, next_class): self._acct = next_class(rec)
        logger.verbose("account.Worker.ensure_created: %s, running=%r"%(self.name,self.is_running()))

        # reservation_alive is set on reervable nodes, and its value is a boolean
        if 'reservation_alive' in rec:
            # reservable nodes
            if rec['reservation_alive']:
                # this sliver has the lease, it is safe to start it
                if not self.is_running(): self.start(rec)
                else: self.configure(rec)
            else:
                # not having the lease, do not start it
                self.configure(rec)
        # usual nodes - preserve old code
        # xxx it's not clear what to do when a sliver changes type/class
        # in a reservable node
        else:
            if not self.is_running() or next_class != curr_class:
                self.start(rec)
            else: self.configure(rec)

    def ensure_destroyed(self): self._destroy(self._get_class())

    def start(self, rec, d = 0):
        self._acct.configure(rec)
        self._acct.start(delay=d)

    def configure(self, rec):
        self._acct.configure(rec)

    def stop(self): self._acct.stop()

    def is_running(self):
        if (self._acct != None) and self._acct.is_running():
            status = True
        else:
            status = False
            logger.verbose("account: Worker(%s): is not running" % self.name)
        return status

    def _destroy(self, curr_class):
        self._acct = None
        if curr_class:
            destroy_sem.acquire()
            try: curr_class.destroy(self.name)
            finally: destroy_sem.release()

    def _get_class(self):
        try: shell = pwd.getpwnam(self.name)[6]
        except KeyError: return None
        return shell_acct_class[shell]

<|endoftext|>


<commit_before>import contextlib
import re
from typing import Iterator
from typing import Mapping
from typing import Optional
from typing import TYPE_CHECKING
from typing import TypeVar
from typing import Union

from sqlalchemy import __version__
from sqlalchemy import inspect
from sqlalchemy import schema
from sqlalchemy import sql
from sqlalchemy import types as sqltypes
from sqlalchemy.engine import url
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.schema import CheckConstraint
from sqlalchemy.schema import Column
from sqlalchemy.schema import ForeignKeyConstraint
from sqlalchemy.sql import visitors
from sqlalchemy.sql.elements import BindParameter
from sqlalchemy.sql.elements import quoted_name
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.sql.visitors import traverse

from . import compat

if TYPE_CHECKING:
    from sqlalchemy import Index
    from sqlalchemy import Table
    from sqlalchemy.engine import Connection
    from sqlalchemy.engine import Dialect
    from sqlalchemy.engine import Transaction
    from sqlalchemy.engine.reflection import Inspector
    from sqlalchemy.sql.base import ColumnCollection
    from sqlalchemy.sql.compiler import SQLCompiler
    from sqlalchemy.sql.dml import Insert
    from sqlalchemy.sql.elements import ColumnClause
    from sqlalchemy.sql.elements import ColumnElement
    from sqlalchemy.sql.schema import Constraint
    from sqlalchemy.sql.schema import SchemaItem
    from sqlalchemy.sql.selectable import Select
    from sqlalchemy.sql.selectable import TableClause

_CE = TypeVar("_CE", bound=Union["ColumnElement", "SchemaItem"])

def _safe_int(value: str) -> Union[int, str]:
    try:
        return int(value)
    except:
        return value

_vers = tuple(
    [_safe_int(x) for x in re.findall(r"(\d+|[abc]\d)", __version__)]
)
sqla_13 = _vers >= (1, 3)
sqla_14 = _vers >= (1, 4)
sqla_14_26 = _vers >= (1, 4, 26)
sqla_1x = _vers < (2,)

try:
    from sqlalchemy import Computed  # noqa
except ImportError:
    Computed = type(None)  # type: ignore
    has_computed = False
    has_computed_reflection = False
else:
    has_computed = True
    has_computed_reflection = _vers >= (1, 3, 16)

try:
    from sqlalchemy import Identity  # noqa
except ImportError:
    Identity = type(None)  # type: ignore
    has_identity = False
else:
    # attributes common to Indentity and Sequence
    _identity_options_attrs = (
        "start",
        "increment",
        "minvalue",
        "maxvalue",
        "nominvalue",
        "nomaxvalue",
        "cycle",
        "cache",
        "order",
    )
    # attributes of Indentity
    _identity_attrs = _identity_options_attrs + ("on_null",)
    has_identity = True

AUTOINCREMENT_DEFAULT = "auto"

@contextlib.contextmanager
def _ensure_scope_for_ddl(
    connection: Optional["Connection"],
) -> Iterator[None]:
    try:
        in_transaction = connection.in_transaction  # type: ignore[union-attr]
    except AttributeError:
        # catch for MockConnection, None
        yield
    else:
        if not in_transaction():
            assert connection is not None
            with connection.begin():
                yield
        else:
            yield

def _safe_begin_connection_transaction(
    connection: "Connection",
) -> "Transaction":
    transaction = _get_connection_transaction(connection)
    if transaction:
        return transaction
    else:
        return connection.begin()

def _safe_commit_connection_transaction(
    connection: "Connection",
) -> None:
    transaction = _get_connection_transaction(connection)
    if transaction:
        transaction.commit()

def _safe_rollback_connection_transaction(
    connection: "Connection",
) -> None:
    transaction = _get_connection_transaction(connection)
    if transaction:
        transaction.rollback()

def _get_connection_in_transaction(connection: Optional["Connection"]) -> bool:
    try:
        in_transaction = connection.in_transaction  # type: ignore
    except AttributeError:
        # catch for MockConnection
        return False
    else:
        return in_transaction()

def _copy(schema_item: _CE, **kw) -> _CE:
    if hasattr(schema_item, "_copy"):
        return schema_item._copy(**kw)
    else:
        return schema_item.copy(**kw)

def _get_connection_transaction(
    connection: "Connection",
) -> Optional["Transaction"]:
    if sqla_14:
        return connection.get_transaction()
    else:
        r = connection._root  # type: ignore[attr-defined]
        return r._Connection__transaction

def _create_url(*arg, **kw) -> url.URL:
    if hasattr(url.URL, "create"):
        return url.URL.create(*arg, **kw)
    else:
        return url.URL(*arg, **kw)

def _connectable_has_table(
    connectable: "Connection", tablename: str, schemaname: Union[str, None]
) -> bool:
    if sqla_14:
        return inspect(connectable).has_table(tablename, schemaname)
    else:
        return connectable.dialect.has_table(
            connectable, tablename, schemaname
        )

def _exec_on_inspector(inspector, statement, **params):
    if sqla_14:
        with inspector._operation_context() as conn:
            return conn.execute(statement, params)
    else:
        return inspector.bind.execute(statement, params)

def _nullability_might_be_unset(metadata_column):
    if not sqla_14:
        return metadata_column.nullable
    else:
        from sqlalchemy.sql import schema

        return (
            metadata_column._user_defined_nullable is schema.NULL_UNSPECIFIED
        )

def _server_default_is_computed(*server_default) -> bool:
    if not has_computed:
        return False
    else:
        return any(isinstance(sd, Computed) for sd in server_default)

def _server_default_is_identity(*server_default) -> bool:
    if not sqla_14:
        return False
    else:
        return any(isinstance(sd, Identity) for sd in server_default)

def _table_for_constraint(constraint: "Constraint") -> "Table":
    if isinstance(constraint, ForeignKeyConstraint):
        table = constraint.parent
        assert table is not None
        return table
    else:
        return constraint.table

def _columns_for_constraint(constraint):
    if isinstance(constraint, ForeignKeyConstraint):
        return [fk.parent for fk in constraint.elements]
    elif isinstance(constraint, CheckConstraint):
        return _find_columns(constraint.sqltext)
    else:
        return list(constraint.columns)

def _reflect_table(
    inspector: "Inspector", table: "Table", include_cols: None
) -> None:
    if sqla_14:
        return inspector.reflect_table(table, None)
    else:
        return inspector.reflecttable(table, None)

def _fk_spec(constraint):
    source_columns = [
        constraint.columns[key].name for key in constraint.column_keys
    ]

    source_table = constraint.parent.name
    source_schema = constraint.parent.schema
    target_schema = constraint.elements[0].column.table.schema
    target_table = constraint.elements[0].column.table.name
    target_columns = [element.column.name for element in constraint.elements]
    ondelete = constraint.ondelete
    onupdate = constraint.onupdate
    deferrable = constraint.deferrable
    initially = constraint.initially
    return (
        source_schema,
        source_table,
        source_columns,
        target_schema,
        target_table,
        target_columns,
        onupdate,
        ondelete,
        deferrable,
        initially,
    )

def _fk_is_self_referential(constraint: "ForeignKeyConstraint") -> bool:
    spec = constraint.elements[0]._get_colspec()  # type: ignore[attr-defined]
    tokens = spec.split(".")
    tokens.pop(-1)  # colname
    tablekey = ".".join(tokens)
    assert constraint.parent is not None
    return tablekey == constraint.parent.key

def _is_type_bound(constraint: "Constraint") -> bool:
    # this deals with SQLAlchemy #3260, don't copy CHECK constraints
    # that will be generated by the type.
    # new feature added for #3260
    return constraint._type_bound  # type: ignore[attr-defined]

def _find_columns(clause):
    """locate Column objects within the given expression."""

    cols = set()
    traverse(clause, {}, {"column": cols.add})
    return cols

def _remove_column_from_collection(
    collection: "ColumnCollection", column: Union["Column", "ColumnClause"]
) -> None:
    """remove a column from a ColumnCollection."""

    # workaround for older SQLAlchemy, remove the
    # same object that's present
    assert column.key is not None
    to_remove = collection[column.key]
    collection.remove(to_remove)

def _textual_index_column(
    table: "Table", text_: Union[str, "TextClause", "ColumnElement"]
) -> Union["ColumnElement", "Column"]:
    """a workaround for the Index construct's severe lack of flexibility"""
    if isinstance(text_, compat.string_types):
        c = Column(text_, sqltypes.NULLTYPE)
        table.append_column(c)
        return c
    elif isinstance(text_, TextClause):
        return _textual_index_element(table, text_)
    elif isinstance(text_, sql.ColumnElement):
        return _copy_expression(text_, table)
    else:
        raise ValueError("String or text() construct expected")

def _copy_expression(expression: _CE, target_table: "Table") -> _CE:
    def replace(col):
        if (
            isinstance(col, Column)
            and col.table is not None
            and col.table is not target_table
        ):
            if col.name in target_table.c:
                return target_table.c[col.name]
            else:
                c = _copy(col)
                target_table.append_column(c)
                return c
        else:
            return None

    return visitors.replacement_traverse(expression, {}, replace)

class _textual_index_element(sql.ColumnElement):
    """Wrap around a sqlalchemy text() construct in such a way that
    we appear like a column-oriented SQL expression to an Index
    construct.

    The issue here is that currently the Postgresql dialect, the biggest
    recipient of functional indexes, keys all the index expressions to
    the corresponding column expressions when rendering CREATE INDEX,
    so the Index we create here needs to have a .columns collection that
    is the same length as the .expressions collection.  Ultimately
    SQLAlchemy should support text() expressions in indexes.

    See SQLAlchemy issue 3174.

    """

    __visit_name__ = "_textual_idx_element"

    def __init__(self, table: "Table", text: "TextClause") -> None:
        self.table = table
        self.text = text
        self.key = text.text
        self.fake_column = schema.Column(self.text.text, sqltypes.NULLTYPE)
        table.append_column(self.fake_column)

    def get_children(self):
        return [self.fake_column]

@compiles(_textual_index_element)
def _render_textual_index_column(
    element: _textual_index_element, compiler: "SQLCompiler", **kw
) -> str:
    return compiler.process(element.text, **kw)

class _literal_bindparam(BindParameter):
    pass

@compiles(_literal_bindparam)
def _render_literal_bindparam(
    element: _literal_bindparam, compiler: "SQLCompiler", **kw
) -> str:
    return compiler.render_literal_bindparam(element, **kw)

def _get_index_expressions(idx):
    return list(idx.expressions)

def _get_index_column_names(idx):
    return [getattr(exp, "name", None) for exp in _get_index_expressions(idx)]

def _column_kwargs(col: "Column") -> Mapping:
    if sqla_13:
        return col.kwargs
    else:
        return {}

def _get_constraint_final_name(
    constraint: Union["Index", "Constraint"], dialect: Optional["Dialect"]
) -> Optional[str]:
    if constraint.name is None:
        return None
    assert dialect is not None
    if sqla_14:
        # for SQLAlchemy 1.4 we would like to have the option to expand
        # the use of "deferred" names for constraints as well as to have
        # some flexibility with "None" name and similar; make use of new
        # SQLAlchemy API to return what would be the final compiled form of
        # the name for this dialect.
        return dialect.identifier_preparer.format_constraint(
            constraint, _alembic_quote=False
        )
    else:

        # prior to SQLAlchemy 1.4, work around quoting logic to get at the
        # final compiled name without quotes.
        if hasattr(constraint.name, "quote"):
            # might be quoted_name, might be truncated_name, keep it the
            # same
            quoted_name_cls: type = type(constraint.name)
        else:
            quoted_name_cls = quoted_name

        new_name = quoted_name_cls(str(constraint.name), quote=False)
        constraint = constraint.__class__(name=new_name)

        if isinstance(constraint, schema.Index):
            # name should not be quoted.
            d = dialect.ddl_compiler(dialect, None)
            return d._prepared_index_name(  # type: ignore[attr-defined]
                constraint
            )
        else:
            # name should not be quoted.
            return dialect.identifier_preparer.format_constraint(constraint)

def _constraint_is_named(
    constraint: Union["Constraint", "Index"], dialect: Optional["Dialect"]
) -> bool:
    if sqla_14:
        if constraint.name is None:
            return False
        assert dialect is not None
        name = dialect.identifier_preparer.format_constraint(
            constraint, _alembic_quote=False
        )
        return name is not None
    else:
        return constraint.name is not None

def _is_mariadb(mysql_dialect: "Dialect") -> bool:
    if sqla_14:
        return mysql_dialect.is_mariadb  # type: ignore[attr-defined]
    else:
        return bool(
            mysql_dialect.server_version_info
            and mysql_dialect._is_mariadb  # type: ignore[attr-defined]
        )

def _mariadb_normalized_version_info(mysql_dialect):
    return mysql_dialect._mariadb_normalized_version_info

def _insert_inline(table: Union["TableClause", "Table"]) -> "Insert":
    if sqla_14:
        return table.insert().inline()
    else:
        return table.insert(inline=True)

if sqla_14:
    from sqlalchemy import create_mock_engine
    from sqlalchemy import select as _select
else:
    from sqlalchemy import create_engine

    def create_mock_engine(url, executor, **kw):  # type: ignore[misc]
        return create_engine(
            "postgresql://", strategy="mock", executor=executor
        )

    def _select(*columns, **kw) -> "Select":
        return sql.select(list(columns), **kw)
<commit_msg>use a heuristic for pure "future engine" until future engine merges<commit_after>import contextlib
import re
from typing import Iterator
from typing import Mapping
from typing import Optional
from typing import TYPE_CHECKING
from typing import TypeVar
from typing import Union

from sqlalchemy import __version__
from sqlalchemy import inspect
from sqlalchemy import schema
from sqlalchemy import sql
from sqlalchemy import types as sqltypes
from sqlalchemy.engine import url
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.schema import CheckConstraint
from sqlalchemy.schema import Column
from sqlalchemy.schema import ForeignKeyConstraint
from sqlalchemy.sql import visitors
from sqlalchemy.sql.elements import BindParameter
from sqlalchemy.sql.elements import quoted_name
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.sql.visitors import traverse

from . import compat

if TYPE_CHECKING:
    from sqlalchemy import Index
    from sqlalchemy import Table
    from sqlalchemy.engine import Connection
    from sqlalchemy.engine import Dialect
    from sqlalchemy.engine import Transaction
    from sqlalchemy.engine.reflection import Inspector
    from sqlalchemy.sql.base import ColumnCollection
    from sqlalchemy.sql.compiler import SQLCompiler
    from sqlalchemy.sql.dml import Insert
    from sqlalchemy.sql.elements import ColumnClause
    from sqlalchemy.sql.elements import ColumnElement
    from sqlalchemy.sql.schema import Constraint
    from sqlalchemy.sql.schema import SchemaItem
    from sqlalchemy.sql.selectable import Select
    from sqlalchemy.sql.selectable import TableClause

_CE = TypeVar("_CE", bound=Union["ColumnElement", "SchemaItem"])

def _safe_int(value: str) -> Union[int, str]:
    try:
        return int(value)
    except:
        return value

_vers = tuple(
    [_safe_int(x) for x in re.findall(r"(\d+|[abc]\d)", __version__)]
)
sqla_13 = _vers >= (1, 3)
sqla_14 = _vers >= (1, 4)
sqla_14_26 = _vers >= (1, 4, 26)

if sqla_14:
    # when future engine merges, this can be again based on version string
    from sqlalchemy.engine import Connection as legacy_connection

    sqla_1x = not hasattr(legacy_connection, "commit")
else:
    sqla_1x = True

try:
    from sqlalchemy import Computed  # noqa
except ImportError:
    Computed = type(None)  # type: ignore
    has_computed = False
    has_computed_reflection = False
else:
    has_computed = True
    has_computed_reflection = _vers >= (1, 3, 16)

try:
    from sqlalchemy import Identity  # noqa
except ImportError:
    Identity = type(None)  # type: ignore
    has_identity = False
else:
    # attributes common to Indentity and Sequence
    _identity_options_attrs = (
        "start",
        "increment",
        "minvalue",
        "maxvalue",
        "nominvalue",
        "nomaxvalue",
        "cycle",
        "cache",
        "order",
    )
    # attributes of Indentity
    _identity_attrs = _identity_options_attrs + ("on_null",)
    has_identity = True

AUTOINCREMENT_DEFAULT = "auto"

@contextlib.contextmanager
def _ensure_scope_for_ddl(
    connection: Optional["Connection"],
) -> Iterator[None]:
    try:
        in_transaction = connection.in_transaction  # type: ignore[union-attr]
    except AttributeError:
        # catch for MockConnection, None
        yield
    else:
        if not in_transaction():
            assert connection is not None
            with connection.begin():
                yield
        else:
            yield

def _safe_begin_connection_transaction(
    connection: "Connection",
) -> "Transaction":
    transaction = _get_connection_transaction(connection)
    if transaction:
        return transaction
    else:
        return connection.begin()

def _safe_commit_connection_transaction(
    connection: "Connection",
) -> None:
    transaction = _get_connection_transaction(connection)
    if transaction:
        transaction.commit()

def _safe_rollback_connection_transaction(
    connection: "Connection",
) -> None:
    transaction = _get_connection_transaction(connection)
    if transaction:
        transaction.rollback()

def _get_connection_in_transaction(connection: Optional["Connection"]) -> bool:
    try:
        in_transaction = connection.in_transaction  # type: ignore
    except AttributeError:
        # catch for MockConnection
        return False
    else:
        return in_transaction()

def _copy(schema_item: _CE, **kw) -> _CE:
    if hasattr(schema_item, "_copy"):
        return schema_item._copy(**kw)
    else:
        return schema_item.copy(**kw)

def _get_connection_transaction(
    connection: "Connection",
) -> Optional["Transaction"]:
    if sqla_14:
        return connection.get_transaction()
    else:
        r = connection._root  # type: ignore[attr-defined]
        return r._Connection__transaction

def _create_url(*arg, **kw) -> url.URL:
    if hasattr(url.URL, "create"):
        return url.URL.create(*arg, **kw)
    else:
        return url.URL(*arg, **kw)

def _connectable_has_table(
    connectable: "Connection", tablename: str, schemaname: Union[str, None]
) -> bool:
    if sqla_14:
        return inspect(connectable).has_table(tablename, schemaname)
    else:
        return connectable.dialect.has_table(
            connectable, tablename, schemaname
        )

def _exec_on_inspector(inspector, statement, **params):
    if sqla_14:
        with inspector._operation_context() as conn:
            return conn.execute(statement, params)
    else:
        return inspector.bind.execute(statement, params)

def _nullability_might_be_unset(metadata_column):
    if not sqla_14:
        return metadata_column.nullable
    else:
        from sqlalchemy.sql import schema

        return (
            metadata_column._user_defined_nullable is schema.NULL_UNSPECIFIED
        )

def _server_default_is_computed(*server_default) -> bool:
    if not has_computed:
        return False
    else:
        return any(isinstance(sd, Computed) for sd in server_default)

def _server_default_is_identity(*server_default) -> bool:
    if not sqla_14:
        return False
    else:
        return any(isinstance(sd, Identity) for sd in server_default)

def _table_for_constraint(constraint: "Constraint") -> "Table":
    if isinstance(constraint, ForeignKeyConstraint):
        table = constraint.parent
        assert table is not None
        return table
    else:
        return constraint.table

def _columns_for_constraint(constraint):
    if isinstance(constraint, ForeignKeyConstraint):
        return [fk.parent for fk in constraint.elements]
    elif isinstance(constraint, CheckConstraint):
        return _find_columns(constraint.sqltext)
    else:
        return list(constraint.columns)

def _reflect_table(
    inspector: "Inspector", table: "Table", include_cols: None
) -> None:
    if sqla_14:
        return inspector.reflect_table(table, None)
    else:
        return inspector.reflecttable(table, None)

def _fk_spec(constraint):
    source_columns = [
        constraint.columns[key].name for key in constraint.column_keys
    ]

    source_table = constraint.parent.name
    source_schema = constraint.parent.schema
    target_schema = constraint.elements[0].column.table.schema
    target_table = constraint.elements[0].column.table.name
    target_columns = [element.column.name for element in constraint.elements]
    ondelete = constraint.ondelete
    onupdate = constraint.onupdate
    deferrable = constraint.deferrable
    initially = constraint.initially
    return (
        source_schema,
        source_table,
        source_columns,
        target_schema,
        target_table,
        target_columns,
        onupdate,
        ondelete,
        deferrable,
        initially,
    )

def _fk_is_self_referential(constraint: "ForeignKeyConstraint") -> bool:
    spec = constraint.elements[0]._get_colspec()  # type: ignore[attr-defined]
    tokens = spec.split(".")
    tokens.pop(-1)  # colname
    tablekey = ".".join(tokens)
    assert constraint.parent is not None
    return tablekey == constraint.parent.key

def _is_type_bound(constraint: "Constraint") -> bool:
    # this deals with SQLAlchemy #3260, don't copy CHECK constraints
    # that will be generated by the type.
    # new feature added for #3260
    return constraint._type_bound  # type: ignore[attr-defined]

def _find_columns(clause):
    """locate Column objects within the given expression."""

    cols = set()
    traverse(clause, {}, {"column": cols.add})
    return cols

def _remove_column_from_collection(
    collection: "ColumnCollection", column: Union["Column", "ColumnClause"]
) -> None:
    """remove a column from a ColumnCollection."""

    # workaround for older SQLAlchemy, remove the
    # same object that's present
    assert column.key is not None
    to_remove = collection[column.key]
    collection.remove(to_remove)

def _textual_index_column(
    table: "Table", text_: Union[str, "TextClause", "ColumnElement"]
) -> Union["ColumnElement", "Column"]:
    """a workaround for the Index construct's severe lack of flexibility"""
    if isinstance(text_, compat.string_types):
        c = Column(text_, sqltypes.NULLTYPE)
        table.append_column(c)
        return c
    elif isinstance(text_, TextClause):
        return _textual_index_element(table, text_)
    elif isinstance(text_, sql.ColumnElement):
        return _copy_expression(text_, table)
    else:
        raise ValueError("String or text() construct expected")

def _copy_expression(expression: _CE, target_table: "Table") -> _CE:
    def replace(col):
        if (
            isinstance(col, Column)
            and col.table is not None
            and col.table is not target_table
        ):
            if col.name in target_table.c:
                return target_table.c[col.name]
            else:
                c = _copy(col)
                target_table.append_column(c)
                return c
        else:
            return None

    return visitors.replacement_traverse(expression, {}, replace)

class _textual_index_element(sql.ColumnElement):
    """Wrap around a sqlalchemy text() construct in such a way that
    we appear like a column-oriented SQL expression to an Index
    construct.

    The issue here is that currently the Postgresql dialect, the biggest
    recipient of functional indexes, keys all the index expressions to
    the corresponding column expressions when rendering CREATE INDEX,
    so the Index we create here needs to have a .columns collection that
    is the same length as the .expressions collection.  Ultimately
    SQLAlchemy should support text() expressions in indexes.

    See SQLAlchemy issue 3174.

    """

    __visit_name__ = "_textual_idx_element"

    def __init__(self, table: "Table", text: "TextClause") -> None:
        self.table = table
        self.text = text
        self.key = text.text
        self.fake_column = schema.Column(self.text.text, sqltypes.NULLTYPE)
        table.append_column(self.fake_column)

    def get_children(self):
        return [self.fake_column]

@compiles(_textual_index_element)
def _render_textual_index_column(
    element: _textual_index_element, compiler: "SQLCompiler", **kw
) -> str:
    return compiler.process(element.text, **kw)

class _literal_bindparam(BindParameter):
    pass

@compiles(_literal_bindparam)
def _render_literal_bindparam(
    element: _literal_bindparam, compiler: "SQLCompiler", **kw
) -> str:
    return compiler.render_literal_bindparam(element, **kw)

def _get_index_expressions(idx):
    return list(idx.expressions)

def _get_index_column_names(idx):
    return [getattr(exp, "name", None) for exp in _get_index_expressions(idx)]

def _column_kwargs(col: "Column") -> Mapping:
    if sqla_13:
        return col.kwargs
    else:
        return {}

def _get_constraint_final_name(
    constraint: Union["Index", "Constraint"], dialect: Optional["Dialect"]
) -> Optional[str]:
    if constraint.name is None:
        return None
    assert dialect is not None
    if sqla_14:
        # for SQLAlchemy 1.4 we would like to have the option to expand
        # the use of "deferred" names for constraints as well as to have
        # some flexibility with "None" name and similar; make use of new
        # SQLAlchemy API to return what would be the final compiled form of
        # the name for this dialect.
        return dialect.identifier_preparer.format_constraint(
            constraint, _alembic_quote=False
        )
    else:

        # prior to SQLAlchemy 1.4, work around quoting logic to get at the
        # final compiled name without quotes.
        if hasattr(constraint.name, "quote"):
            # might be quoted_name, might be truncated_name, keep it the
            # same
            quoted_name_cls: type = type(constraint.name)
        else:
            quoted_name_cls = quoted_name

        new_name = quoted_name_cls(str(constraint.name), quote=False)
        constraint = constraint.__class__(name=new_name)

        if isinstance(constraint, schema.Index):
            # name should not be quoted.
            d = dialect.ddl_compiler(dialect, None)
            return d._prepared_index_name(  # type: ignore[attr-defined]
                constraint
            )
        else:
            # name should not be quoted.
            return dialect.identifier_preparer.format_constraint(constraint)

def _constraint_is_named(
    constraint: Union["Constraint", "Index"], dialect: Optional["Dialect"]
) -> bool:
    if sqla_14:
        if constraint.name is None:
            return False
        assert dialect is not None
        name = dialect.identifier_preparer.format_constraint(
            constraint, _alembic_quote=False
        )
        return name is not None
    else:
        return constraint.name is not None

def _is_mariadb(mysql_dialect: "Dialect") -> bool:
    if sqla_14:
        return mysql_dialect.is_mariadb  # type: ignore[attr-defined]
    else:
        return bool(
            mysql_dialect.server_version_info
            and mysql_dialect._is_mariadb  # type: ignore[attr-defined]
        )

def _mariadb_normalized_version_info(mysql_dialect):
    return mysql_dialect._mariadb_normalized_version_info

def _insert_inline(table: Union["TableClause", "Table"]) -> "Insert":
    if sqla_14:
        return table.insert().inline()
    else:
        return table.insert(inline=True)

if sqla_14:
    from sqlalchemy import create_mock_engine
    from sqlalchemy import select as _select
else:
    from sqlalchemy import create_engine

    def create_mock_engine(url, executor, **kw):  # type: ignore[misc]
        return create_engine(
            "postgresql://", strategy="mock", executor=executor
        )

    def _select(*columns, **kw) -> "Select":
        return sql.select(list(columns), **kw)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
from core import *
import wi

class push(Command):
    """Push pending commits to TFS."""

    def initArgParser(self, parser):
        Command.initArgParser(self, parser)
        parser.add_argument('--number',
            help='maximum number of changesets to push',
            type=int,
            default=None)

    def _push(self, hash, index, total):
        dryRun = self.args.dryRun
        print()
        printLine()
        git('checkout ' + hash)
        print('Pushing [%d/%d] %s...' % (index + 1, total, git(r'log -1 --format="%h \"%s\""')))

        def rawDiff(changeType):
            return git('diff --raw --find-copies-harder HEAD^.. --diff-filter=%s' % changeType)

        def readChanges(changeType, displayChangeType):
            changes = [change[change.index('\t'):].strip().split('\t') for change in rawDiff(changeType).splitlines()]
            if changes:
                print(displayChangeType + ':')
                indentPrint([' -> '.join(f) for f in changes])
                yield changes

        def joinFiles(files):
            return '"' + '" "'.join(files) + '"'
        def joinChanges(changes):
            return ' '.join(map(joinFiles, changes))

        unknownChanges = rawDiff('TUX')
        if unknownChanges:
            print('Unexpected file change!!')
            print()
            indentPrint(unknownChanges)
            fail()

        def tfmut(args):
            tf(args, dryRun = dryRun)

        try:
            for c in readChanges('D', 'Removed'):
                tfmut('rm -recursive ' + joinChanges(c))
            for c in readChanges('M', 'Modified'):
                tfmut('checkout ' + joinChanges(c))
            for changes in readChanges('R', 'Renamed'):
                for files in changes:
                    src, dest = files
                    if not dryRun:
                        os.rename(dest, src)
                    try:
                        tfmut('rename ' + joinFiles(files))
                    except:
                        if not dryRun:
                            os.rename(src, dest)
                        raise
            for c in readChanges('CA', 'Added'):
                tfmut('add ' + joinChanges([files[-1:] for files in c]))

            print('Checking in...')
            comment = git('log -1 --format=%s%n%b').strip()
            workitems = git('notes --ref=%s show %s' % (wi.noteNamespace, hash), errorValue='')
            if workitems:
                workitems = '"-associate:%s"' % workitems
            checkin = tf('checkin "-comment:%s" -recursive %s .' % (comment, workitems), output = True, dryRun = dryRun and 'Changeset #12345')
        except:
            if not dryRun:
                print('Restoring Git and TFS state...')
                with ReadOnlyWorktree():
                    tf('undo -recursive .', allowedExitCodes = [0, 100])
                git('checkout -f tfs')
            raise
        changeSetNumber = re.search(r'^Changeset #(\d+)', checkin, re.M).group(1)

        # add a note about the changeset number
        print('Moving tfs branch HEAD and marking the commit with a "tf" note')
        git('checkout tfs', dryRun = dryRun)
        git('merge --ff-only %s' % hash, dryRun = dryRun)
        git('notes add -m "%s" %s' % (changeSetNumber, hash), dryRun = dryRun)

    def _run(self):
        print('Pushing to TFS')
        lastCommit = git('log -1 --format=%H tfs')
        lastMasterCommit = git('log master --format=%H -1')
        unmergedCommits = git('log %s.. master --oneline' % lastMasterCommit)
        if unmergedCommits:
            print('You have unmerged changes in tfs branch:')
            indentPrint(unmergedCommits)
            fail()

        print('Last synchronized commit:', git('log -1 --format=%h tfs'))
        commits = git('log %s.. --format=%%H master --reverse --first-parent' % lastCommit).splitlines()
        commits = commits[:self.args.number]
        if not commits:
            print('Nothing to push')
            return

        print('Checking whether there are no unfetched changes on TFS...')
        latestGitChangeset = re.findall(r'^\d+', git('notes show ' + lastCommit), re.M)[-1]
        latestTfChangeset = tf.history('-stopafter:1')[0].id
        if int(latestGitChangeset) < int(latestTfChangeset):
            print('There are unfetched changes on TFS. Fetch and merge them before pushing')
            print('Latest local changeset:', latestGitChangeset)
            print('Latest TFS changeset:', latestTfChangeset)
            fail()

        print('%d commit(s) to be pushed:' % len(commits))

        for i, hash in enumerate(commits):
            self._push(hash, i, len(commits))

if __name__ == '__main__':
    push().run()<commit_msg>'tf checkin' is allowed to exit with code 1<commit_after>#!/usr/bin/env python3
from core import *
import wi

class push(Command):
    """Push pending commits to TFS."""

    def initArgParser(self, parser):
        Command.initArgParser(self, parser)
        parser.add_argument('--number',
            help='maximum number of changesets to push',
            type=int,
            default=None)

    def _push(self, hash, index, total):
        dryRun = self.args.dryRun
        print()
        printLine()
        git('checkout ' + hash)
        print('Pushing [%d/%d] %s...' % (index + 1, total, git(r'log -1 --format="%h \"%s\""')))

        def rawDiff(changeType):
            return git('diff --raw --find-copies-harder HEAD^.. --diff-filter=%s' % changeType)

        def readChanges(changeType, displayChangeType):
            changes = [change[change.index('\t'):].strip().split('\t') for change in rawDiff(changeType).splitlines()]
            if changes:
                print(displayChangeType + ':')
                indentPrint([' -> '.join(f) for f in changes])
                yield changes

        def joinFiles(files):
            return '"' + '" "'.join(files) + '"'
        def joinChanges(changes):
            return ' '.join(map(joinFiles, changes))

        unknownChanges = rawDiff('TUX')
        if unknownChanges:
            print('Unexpected file change!!')
            print()
            indentPrint(unknownChanges)
            fail()

        def tfmut(args):
            tf(args, dryRun = dryRun)

        try:
            for c in readChanges('D', 'Removed'):
                tfmut('rm -recursive ' + joinChanges(c))
            for c in readChanges('M', 'Modified'):
                tfmut('checkout ' + joinChanges(c))
            for changes in readChanges('R', 'Renamed'):
                for files in changes:
                    src, dest = files
                    if not dryRun:
                        os.rename(dest, src)
                    try:
                        tfmut('rename ' + joinFiles(files))
                    except:
                        if not dryRun:
                            os.rename(src, dest)
                        raise
            for c in readChanges('CA', 'Added'):
                tfmut('add ' + joinChanges([files[-1:] for files in c]))

            print('Checking in...')
            comment = git('log -1 --format=%s%n%b').strip()
            workitems = git('notes --ref=%s show %s' % (wi.noteNamespace, hash), errorValue='')
            if workitems:
                workitems = '"-associate:%s"' % workitems
            checkin = tf('checkin "-comment:%s" -recursive %s .' % (comment, workitems),
                allowedExitCodes=[0, 1],
                output=True,
                dryRun=dryRun and 'Changeset #12345')
            changeSetNumber = re.search(r'^Changeset #(\d+)', checkin, re.M)
            if not changeSetNumber:
                fail('Check in failed.')
            changeSetNumber = changeSetNumber.group(1)
        except:
            if not dryRun:
                print('Restoring Git and TFS state...')
                with ReadOnlyWorktree():
                    tf('undo -recursive .', allowedExitCodes = [0, 100])
                git('checkout -f tfs')
            raise

        # add a note about the changeset number
        print('Moving tfs branch HEAD and marking the commit with a "tf" note')
        git('checkout tfs', dryRun = dryRun)
        git('merge --ff-only %s' % hash, dryRun = dryRun)
        git('notes add -m "%s" %s' % (changeSetNumber, hash), dryRun = dryRun)

    def _run(self):
        print('Pushing to TFS')
        lastCommit = git('log -1 --format=%H tfs')
        lastMasterCommit = git('log master --format=%H -1')
        unmergedCommits = git('log %s.. master --oneline' % lastMasterCommit)
        if unmergedCommits:
            print('You have unmerged changes in tfs branch:')
            indentPrint(unmergedCommits)
            fail()

        print('Last synchronized commit:', git('log -1 --format=%h tfs'))
        commits = git('log %s.. --format=%%H master --reverse --first-parent' % lastCommit).splitlines()
        commits = commits[:self.args.number]
        if not commits:
            print('Nothing to push')
            return

        print('Checking whether there are no unfetched changes on TFS...')
        latestGitChangeset = re.findall(r'^\d+', git('notes show ' + lastCommit), re.M)[-1]
        latestTfChangeset = tf.history('-stopafter:1')[0].id
        if int(latestGitChangeset) < int(latestTfChangeset):
            print('There are unfetched changes on TFS. Fetch and merge them before pushing')
            print('Latest local changeset:', latestGitChangeset)
            print('Latest TFS changeset:', latestTfChangeset)
            fail()

        print('%d commit(s) to be pushed:' % len(commits))

        for i, hash in enumerate(commits):
            self._push(hash, i, len(commits))

if __name__ == '__main__':
    push().run()<|endoftext|>