<commit_before>import rospy
import yaml
import tf
import tf_conversions.posemath as pm

from predicator_msgs.srv import *
from predicator_msgs.msg import *

from librarian_msgs.srv import *
from librarian_msgs.msg import *

from costar_objrec_msgs.msg import *

from predicator_landmark import GetWaypointsService

'''
SmartWaypointManager
This class will create and load smart waypoints from the $LIBRARIAN_HOME/smart_waypoints directory
It also provides ways to use and access these methods
- stores smartmove waypoints with associated class info
- stores joint space waypoints
- stores a set of 6dof cartesian waypoints (for now in the world frame)
'''
class SmartWaypointManager:

    def __init__(self,world="world",ns="",endpoint="/endpoint", listener = None):
        self.get_waypoints_srv = GetWaypointsService(world=world,service=False)

        rospy.loginfo("[SmartMove] Waiting for LIBRARIAN to handle file I/O...")
        rospy.wait_for_service('/librarian/add_type',5)
        self.add_type_service = rospy.ServiceProxy('/librarian/add_type', librarian_msgs.srv.AddType)
        self.save_service = rospy.ServiceProxy('/librarian/save', librarian_msgs.srv.Save)
        self.load_service = rospy.ServiceProxy('/librarian/load', librarian_msgs.srv.Load)
        self.list_service = rospy.ServiceProxy('/librarian/list', librarian_msgs.srv.List)
        self.delete_service = rospy.ServiceProxy('/librarian/delete', librarian_msgs.srv.Delete)

        self.get_assignment_service = rospy.ServiceProxy('/predicator/get_assignment', predicator_msgs.srv.GetAssignment)

        self.detected_objects = rospy.Subscriber(ns + '/detected_object_list', DetectedObjectList, self.detected_objects_cb)

        self.broadcaster = tf.TransformBroadcaster()
        if not listener is None:
          self.listener = listener
        else:
          self.listener = tf.TransformListener()

        self.world = world
        self.endpoint = endpoint

        self.folder = 'smartmove_waypoint'
        self.js_folder = 'joint_states'
        self.cart_folder = 'cartesian'
        self.info_folder = 'smartmove_info'

        self.add_type_service(self.folder)

        self.waypoints = {}
        self.waypoint_names = {}
        self.js_waypoints = {}
        self.js_waypoint_names = {}
        self.cart_waypoints = {}
        self.cart_waypoint_names = {}

        self.all_moves = []
        self.all_js_moves = []

        self.objs = []
        self.obj_classes = []
        self.obj_class = {}

        self.available_object_classes = None
        self.available_regions = None
        self.available_references = None

        self.add_type_service(self.info_folder)
        self.add_type_service(self.js_folder)
        self.available_obj_classes = yaml.load(self.load_service(type="smartmove_info",id="obj_classes").text)
        self.available_regions = yaml.load(self.load_service(type="smartmove_info",id="regions").text)
        self.available_references = yaml.load(self.load_service(type="smartmove_info",id="references").text)
        rospy.logwarn("[SMARTMOVE] Available classes = " + str(self.available_obj_classes))

    '''
    reads in costar object detection messages
    this includes symmetry information produced by the vision pipeline
    '''
    def detected_objects_cb(self,msg):
        #self.objs = [obj.id for obj in msg.objects]
        #self.obj_classes = [obj.object_class for obj in msg.objects]
        #for (o,c) in zip(self.objs,self.obj_classes):
        #    self.obj_class[o] = c

	for obj in msg.objects:
	  if not obj.id in self.objs:
	    self.objs.append(obj.id)
	  if not obj.object_class in self.obj_classes:
            self.objs.append(obj.object_class)
	  self.obj_class[obj.id] = obj.object_class

        #print self.objs
        #print self.obj_classes

    '''
    get all waypoints from the disk
    now including some joint space waypoints
    '''
    def load_all(self):

        self.waypoints = {}
        self.waypoint_names = {}
        self.all_moves = []

        '''
        this section loads class and pose information for computing smartmoves
        '''

        waypoint_filenames = self.list_service(self.folder).entries

        for name in waypoint_filenames:
          data = yaml.load(self.load_service(id=name,type=self.folder).text)
          if not data[1] in self.waypoints.keys():
            self.waypoints[data[1]] = []
            self.waypoint_names[data[1]] = []

          self.waypoints[data[1]].append(data[0])
          self.waypoint_names[data[1]].append(name)
          self.all_moves.append(data[1] + "/" + name)

        '''
        this section loads joint space waypoints
        '''

        js_filenames = self.list_service(self.js_folder).entries

        for name in js_filenames:
          data = yaml.load(self.load_service(id=name,type=self.js_folder).text)
          if not data[1] in self.waypoints.keys():
            self.js_waypoints[data[1]] = []
            self.js_waypoint_names[data[1]] = []

            self.js_waypoints[data[1]].append(data[0])
            self.js_waypoint_names[data[1]].append(name)
            self.all_js_moves.append(data[1] + "/" + name)

        print " === LOADING === "
        print self.waypoint_names
        print self.waypoints
        print self.js_waypoint_names
        print self.js_waypoint

        '''
        this section loads cartesian waypoints
        '''

        cart_filenames = self.list_service(self.cart_folder).entries
        for name in cart_filenames:
          data = yaml.load(self.load_service(id=name,type=self.cart_folder).text)

    def lookup_waypoint(self,obj_class,name):
      rospy.logwarn("looking for %s"%name)
      rospy.logwarn(self.waypoints)
      rospy.logwarn(self.waypoint_names)
      return self.waypoints[obj_class][self.waypoint_names[obj_class].index(name)]

    def get_reference_frames(self):
      return self.available_references

    def get_available_predicates(self):
      return self.available_regions

    def get_all_moves(self):
        return self.all_moves

    def get_detected_objects(self):
        self.objs = []
        self.obj_classes = []
        rospy.logwarn(str(self.available_obj_classes))
        if not self.available_obj_classes is None:
            rospy.logwarn("- not none")
            for oc in self.available_obj_classes:
                rospy.logwarn("-- %s"%oc)
                resp = self.get_assignment_service(PredicateStatement(predicate=oc,params=["*","",""]))
                oc_objs = [p.params[0] for p in resp.values]
                if len(oc_objs) > 0:
                  self.obj_classes.append(oc)
                for obj in oc_objs:
                    self.obj_class[obj] = oc
                    rospy.logwarn("%s = %s"%(obj,oc))
                self.objs += oc_objs

        return self.objs

    def delete(self,move):
        self.delete_service(id=move.strip('/'),type=self.folder)

    def get_detected_object_classes(self):
        self.get_detected_objects()
        return self.obj_classes

    def get_available_object_classes(self):
        return self.available_object_classes

    def save_new_waypoint(self,obj,name):
        pose = self.get_new_waypoint(obj)
        if not pose is None:
            self.save_service(id=name.strip('/'),type=self.folder,text=yaml.dump([pose,self.obj_class[obj]]))

    def save_joint_states(self,q,name):
        self.save_service(id=name.strip('/'),type=self.js_folder,text=yaml.dump(q))

    def get_new_waypoint(self,obj):
        try:
            (trans,rot) = self.listener.lookupTransform(obj,self.endpoint,rospy.Time(0))
            return pm.toMsg(pm.fromTf((trans,rot)))
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException):
            rospy.logerr('Failed to lookup transform from %s to %s'%(obj,self.endpoint))
        return None

    def get_moves_for_class(self,obj_class):
        wpts = []
        try:
            wpts = self.waypoint_names[obj_class]
        except KeyError:
            wpts = []
        return wpts

    def get_moves_for_object(self,obj):
        try:
            print "obj class = %s"%self.obj_class[obj]
        except KeyError:
            rospy.logerr("Could not find object %s! Are you sure this is a valid object?"%obj)

        wpts = []
        try:
            wpts = self.waypoint_names[self.obj_class[obj]]
        except KeyError:
            wpts = []
        return wpts

    def update_tf(self):

        for key in self.waypoints.keys():
            (poses,names) = self.get_waypoints_srv.get_waypoints(key,[],self.waypoints[key],self.waypoint_names[key])
            for (pose,name) in zip(poses,names):
                (trans,rot) = pm.toTf(pm.fromMsg(pose))
                self.broadcaster.sendTransform(trans,rot,rospy.Time.now(),name,self.world)
                #print (trans, rot, name)

<commit_msg>adding more functionality to the costar smart waypoint manager<commit_after>import rospy
import yaml
import tf
import tf_conversions.posemath as pm

from predicator_msgs.srv import *
from predicator_msgs.msg import *

from librarian_msgs.srv import *
from librarian_msgs.msg import *

from costar_objrec_msgs.msg import *
from costar_robot_msgs.srv import *

from predicator_landmark import GetWaypointsService

'''
SmartWaypointManager
This class will create and load smart waypoints from the $LIBRARIAN_HOME/smart_waypoints directory
It also provides ways to use and access these methods
- stores smartmove waypoints with associated class info
- stores joint space waypoints
- stores a set of 6dof cartesian waypoints (for now in the world frame)
'''
class SmartWaypointManager:

    def __init__(self,world="world",ns="",endpoint="/endpoint", listener = None):
        self.get_waypoints_srv = GetWaypointsService(world=world,service=False)

        rospy.loginfo("[SmartMove] Waiting for LIBRARIAN to handle file I/O...")
        rospy.wait_for_service('/librarian/add_type',5)
        self.add_type_service = rospy.ServiceProxy('/librarian/add_type', librarian_msgs.srv.AddType)
        self.save_service = rospy.ServiceProxy('/librarian/save', librarian_msgs.srv.Save)
        self.load_service = rospy.ServiceProxy('/librarian/load', librarian_msgs.srv.Load)
        self.list_service = rospy.ServiceProxy('/librarian/list', librarian_msgs.srv.List)
        self.delete_service = rospy.ServiceProxy('/librarian/delete', librarian_msgs.srv.Delete)

        self.get_assignment_service = rospy.ServiceProxy('/predicator/get_assignment', predicator_msgs.srv.GetAssignment)

        self.detected_objects = rospy.Subscriber(ns + '/detected_object_list', DetectedObjectList, self.detected_objects_cb)

        self.broadcaster = tf.TransformBroadcaster()
        if not listener is None:
          self.listener = listener
        else:
          self.listener = tf.TransformListener()

        self.world = world
        self.endpoint = endpoint

        self.folder = 'smartmove_waypoint'
        self.js_folder = 'joint_states'
        self.cart_folder = 'cartesian'
        self.info_folder = 'smartmove_info'

        self.add_type_service(self.folder)

        self.waypoints = {}
        self.waypoint_names = {}
        self.js_waypoints = {}
        self.js_waypoint_names = {}
        self.cart_waypoints = {}
        self.cart_waypoint_names = {}

        self.all_moves = []
        self.all_js_moves = []

        self.objs = []
        self.obj_classes = []
        self.obj_class = {}

        self.available_object_classes = None
        self.available_regions = None
        self.available_references = None

        self.add_type_service(self.info_folder)
        self.add_type_service(self.js_folder)
        self.available_obj_classes = yaml.load(self.load_service(type="smartmove_info",id="obj_classes").text)
        self.available_regions = yaml.load(self.load_service(type="smartmove_info",id="regions").text)
        self.available_references = yaml.load(self.load_service(type="smartmove_info",id="references").text)
        rospy.logwarn("[SMARTMOVE] Available classes = " + str(self.available_obj_classes))

    '''
    reads in costar object detection messages
    this includes symmetry information produced by the vision pipeline
    '''
    def detected_objects_cb(self,msg):
        #self.objs = [obj.id for obj in msg.objects]
        #self.obj_classes = [obj.object_class for obj in msg.objects]
        #for (o,c) in zip(self.objs,self.obj_classes):
        #    self.obj_class[o] = c

	for obj in msg.objects:
	  if not obj.id in self.objs:
	    self.objs.append(obj.id)
	  if not obj.object_class in self.obj_classes:
            self.objs.append(obj.object_class)
	  self.obj_class[obj.id] = obj.object_class

        #print self.objs
        #print self.obj_classes

    '''
    get all waypoints from the disk
    now including some joint space waypoints
    '''
    def load_all(self):

        self.waypoints = {}
        self.waypoint_names = {}
        self.all_moves = []

        '''
        this section loads class and pose information for computing smartmoves
        '''

        waypoint_filenames = self.list_service(self.folder).entries

        for name in waypoint_filenames:
          data = yaml.load(self.load_service(id=name,type=self.folder).text)
          if not data[1] in self.waypoints.keys():
            self.waypoints[data[1]] = []
            self.waypoint_names[data[1]] = []

          self.waypoints[data[1]].append(data[0])
          self.waypoint_names[data[1]].append(name)
          self.all_moves.append(data[1] + "/" + name)

        '''
        this section loads joint space waypoints
        '''

        js_filenames = self.list_service(self.js_folder).entries

        for name in js_filenames:
          data = yaml.load(self.load_service(id=name,type=self.js_folder).text)
          if not data[1] in self.waypoints.keys():
            self.js_waypoints[data[1]] = []
            self.js_waypoint_names[data[1]] = []

            self.js_waypoints[data[1]].append(data[0])
            self.js_waypoint_names[data[1]].append(name)
            self.all_js_moves.append(data[1] + "/" + name)

        print " === LOADING === "
        print self.waypoint_names
        print self.waypoints
        print self.js_waypoint_names
        print self.js_waypoint

        '''
        this section loads cartesian waypoints
        '''

        cart_filenames = self.list_service(self.cart_folder).entries
        for name in cart_filenames:
          data = yaml.load(self.load_service(id=name,type=self.cart_folder).text)

    def lookup_waypoint(self,obj_class,name):
      rospy.logwarn("looking for %s"%name)
      rospy.logwarn(self.waypoints)
      rospy.logwarn(self.waypoint_names)
      return self.waypoints[obj_class][self.waypoint_names[obj_class].index(name)]

    def get_reference_frames(self):
      return self.available_references

    def get_available_predicates(self):
      return self.available_regions

    def get_all_moves(self):
        return self.all_moves

    def get_detected_objects(self):
        self.objs = []
        self.obj_classes = []
        rospy.logwarn(str(self.available_obj_classes))
        if not self.available_obj_classes is None:
            rospy.logwarn("- not none")
            for oc in self.available_obj_classes:
                rospy.logwarn("-- %s"%oc)
                resp = self.get_assignment_service(PredicateStatement(predicate=oc,params=["*","",""]))
                oc_objs = [p.params[0] for p in resp.values]
                if len(oc_objs) > 0:
                  self.obj_classes.append(oc)
                for obj in oc_objs:
                    self.obj_class[obj] = oc
                    rospy.logwarn("%s = %s"%(obj,oc))
                self.objs += oc_objs

        return self.objs

    def delete(self,move):
        self.delete_service(id=move.strip('/'),type=self.folder)

    def get_detected_object_classes(self):
        self.get_detected_objects()
        return self.obj_classes

    def get_available_object_classes(self):
        return self.available_object_classes

    def save_new_waypoint(self,obj,name):
        pose = self.get_new_waypoint(obj)
        if not pose is None:
            self.save_service(id=name.strip('/'),type=self.folder,text=yaml.dump([pose,self.obj_class[obj]]))

    def save_joint_states(self,q,name):
        self.save_service(id=name.strip('/'),type=self.js_folder,text=yaml.dump(q))

    def get_new_waypoint(self,obj):
        try:
            (trans,rot) = self.listener.lookupTransform(obj,self.endpoint,rospy.Time(0))
            return pm.toMsg(pm.fromTf((trans,rot)))
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException):
            rospy.logerr('Failed to lookup transform from %s to %s'%(obj,self.endpoint))
        return None

    def get_moves_for_class(self,obj_class):
        wpts = []
        try:
            wpts = self.waypoint_names[obj_class]
        except KeyError:
            wpts = []
        return wpts

    def get_moves_for_object(self,obj):
        try:
            print "obj class = %s"%self.obj_class[obj]
        except KeyError:
            rospy.logerr("Could not find object %s! Are you sure this is a valid object?"%obj)

        wpts = []
        try:
            wpts = self.waypoint_names[self.obj_class[obj]]
        except KeyError:
            wpts = []
        return wpts

    def update_tf(self):

        for key in self.waypoints.keys():
            (poses,names) = self.get_waypoints_srv.get_waypoints(key,[],self.waypoints[key],self.waypoint_names[key])
            for (pose,name) in zip(poses,names):
                (trans,rot) = pm.toTf(pm.fromMsg(pose))
                self.broadcaster.sendTransform(trans,rot,rospy.Time.now(),name,self.world)
                #print (trans, rot, name)

    '''
    Send any non-smart waypoints we might be managing.
    '''
    def publish_cartesian_waypoints(self):

        for (name,pose) in self.cart_waypoints:
            (trans,rot) = pm.toTf(pm.fromMsg(pose))
            self.broadcaster.sendTransform(trans,rot,rospy.Time.now(),name,self.world)

    '''
    implement service to get joint states
    '''
    def get_joint_states_by_name_srv(self, req):
        if req.name in self.js_waypoints
            msg = LookupJointStatesResponse(joint_states=self.js_waypoints[name], ack='SUCCESS')
        else:
            msg = LookupJointStatesResponse(ack='FAILURE - %s not found'%req.name)
        return msg
<|endoftext|>


<commit_before>#!/usr/bin/python

from __future__ import unicode_literals
import csv

import click

import prettytable

def interact(seq, desc, count, row, mtx):
    # status output
    keys = sorted(mtx)
    table = prettytable.PrettyTable(keys)
    table.add_row([
        '\033[92m\u2713\u2713\033[0m' if mtx[key] else '--'
        for key in keys
    ])

    print """
Sequence:    {seq}
    {desc}
Interpreted: {i_balls}-{i_strikes}
Final:       {balls}-{strikes}

{mtx}
""".format(seq=seq, desc=', '.join(desc),
           i_balls=count['balls'], i_strikes=count['strikes'],
           balls=row['BALLS_CT'], strikes=row['STRIKES_CT'],
           mtx=table)

    # keep going?
    if raw_input('Continue? [Y/n]: ').lower() == 'n':
        exit()

@click.command()
@click.option('files', '-f', type=click.File('r'), multiple=True,
                required=True, help='Chadwick-parsed event log')
@click.option('imode', '-i', is_flag=True, default=False,
              help='Enable interactive mode')
@click.option('output', '-o', type=click.File('w'), required=True,
              help='Path to output file')
def extract_sequences(files, imode, output):
    """Converts a pitch sequences into a matrix
    of ball-strike counts.
    """

    mtx = {
        '00': 1, '01': 0, '02': 0,
        '10': 0, '20': 0, '30': 0,
        '11': 0, '12': 0, '21': 0,
        '22': 0, '31': 0, '32': 0,
    }

    # write header row
    writer = csv.writer(output)
    headers = ['pitch_seq_tx']
    headers.extend(['c' + key for key in sorted(mtx)])
    writer.writerow(headers)

    for f in files:
        reader = csv.DictReader(f)

        for row in reader:
            count = {'balls': 0, 'strikes': 0}
            desc = []
            seq = row['PITCH_SEQ_TX']

            for pitch in seq:
                # handle walks
                if pitch in ('B', 'I', 'V'):
                    if count['balls'] < 3:
                        count['balls'] += 1
                        desc.append('Ball')
                    else:
                        # ball 4, take your base
                        break

                # handle fouls that do not terminate an at-bat
                elif pitch in ('F', 'R'):
                    # keep going forever
                    if count['strikes'] < 2:
                        count['strikes'] += 1
                    desc.append('Foul')

                # handle strikes that can terminate an at-bat
                elif pitch in ('C', 'K', 'L', 'O', 'Q', 'S', 'T'):
                    if count['strikes'] == 2:
                        # batter struck out
                        break
                    else:
                        count['strikes'] += 1
                        desc.append('Strike')
                else:
                    continue

                key = '%s%s' % (count['balls'], count['strikes'])
                mtx[key] += 1

            if imode is True:
                interact(seq, desc, count, row, mtx)

            # write output
            values = [seq]
            for key in sorted(mtx):
                values.append(mtx[key])
            writer.writerow(values)

            # reset matrix
            for key in mtx:
                mtx[key] = 0
            mtx['00'] = 1

if __name__ == '__main__':
    extract_sequences()
<commit_msg>emit only unique pitch sequences<commit_after>#!/usr/bin/python

from __future__ import unicode_literals
import csv

import click

import prettytable

def interact(seq, desc, count, row, mtx):
    # status output
    keys = sorted(mtx)
    table = prettytable.PrettyTable(keys)
    table.add_row([
        '\033[92m\u2713\u2713\033[0m' if mtx[key] else '--'
        for key in keys
    ])

    print """
Sequence:    {seq}
    {desc}
Interpreted: {i_balls}-{i_strikes}
Final:       {balls}-{strikes}

{mtx}
""".format(seq=seq, desc=', '.join(desc),
           i_balls=count['balls'], i_strikes=count['strikes'],
           balls=row['BALLS_CT'], strikes=row['STRIKES_CT'],
           mtx=table)

    # keep going?
    if raw_input('Continue? [Y/n]: ').lower() == 'n':
        exit()

@click.command()
@click.option('files', '-f', type=click.File('r'), multiple=True,
                required=True, help='Chadwick-parsed event log')
@click.option('imode', '-i', is_flag=True, default=False,
              help='Enable interactive mode')
@click.option('output', '-o', type=click.File('w'), required=True,
              help='Path to output file')
def extract_sequences(files, imode, output):
    """Converts a pitch sequences into a matrix
    of ball-strike counts.
    """

    seq_obv = set()
    mtx = {
        '00': 1, '01': 0, '02': 0,
        '10': 0, '20': 0, '30': 0,
        '11': 0, '12': 0, '21': 0,
        '22': 0, '31': 0, '32': 0,
    }

    # write header row
    writer = csv.writer(output)
    headers = ['pitch_seq_tx']
    headers.extend(['c' + key for key in sorted(mtx)])
    writer.writerow(headers)

    for f in files:
        reader = csv.DictReader(f)

        for row in reader:
            seq = row['PITCH_SEQ_TX']

            if seq in seq_obv:
                continue

            count = {'balls': 0, 'strikes': 0}
            desc = []

            for pitch in seq:
                # handle walks
                if pitch in ('B', 'I', 'V'):
                    if count['balls'] < 3:
                        count['balls'] += 1
                        desc.append('Ball')
                    else:
                        # ball 4, take your base
                        break

                # handle fouls that do not terminate an at-bat
                elif pitch in ('F', 'R'):
                    # keep going forever
                    if count['strikes'] < 2:
                        count['strikes'] += 1
                    desc.append('Foul')

                # handle strikes that can terminate an at-bat
                elif pitch in ('C', 'K', 'L', 'O', 'Q', 'S', 'T'):
                    if count['strikes'] == 2:
                        # batter struck out
                        break
                    else:
                        count['strikes'] += 1
                        desc.append('Strike')
                else:
                    continue

                key = '%s%s' % (count['balls'], count['strikes'])
                if mtx[key] < 1:
                    mtx[key] += 1

            if imode is True:
                interact(seq, desc, count, row, mtx)

            # write output
            values = [seq]
            for key in sorted(mtx):
                values.append(mtx[key])
            writer.writerow(values)

            # mark sequence as observed
            seq_obv.add(seq)

            # reset matrix
            for key in mtx:
                mtx[key] = 0
            mtx['00'] = 1

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


<commit_before>from django import forms

class OrderForm(forms.Form):
    name = forms.CharField(label='Naam', max_length=200)
    wiebetaaltwat = forms.BooleanField(
        label='Betaalt via Wiebetaaltwat.nl',
        required=False)
<commit_msg>made wiebetaaltwat default to true<commit_after>from django import forms

class OrderForm(forms.Form):
    name = forms.CharField(label='Naam', max_length=200)
    wiebetaaltwat = forms.BooleanField(
        label='Betaalt via Wiebetaaltwat.nl',
        required=False,
        initial=True)
<|endoftext|>


<commit_before>commands = {
    'app': {
        'label': 'Application',
        'actions': {
            'neweditor': {
                'label': 'New SQL editor',
                'description': 'Open new SQL editor',
                'icon': 'document-new-symbolic',
                'shortcut': '<Control>N',
                'callback': 'win.docview.add_worksheet'
            },
            'switch_editor1': {
                'label': 'Switch to editor 1',
                'shortcut': '<Control>1',
                'callback': 'win.docview.switch_to_editor',
                'args': [1]
            },
            'switch_editor2': {
                'label': 'Switch to editor 2',
                'shortcut': '<Control>2',
                'callback': 'win.docview.switch_to_editor',
                'args': [2]
            },
            'switch_editor3': {
                'label': 'Switch to editor 3',
                'shortcut': '<Control>3',
                'callback': 'win.docview.switch_to_editor',
                'args': [3]
            },
            'switch_editor4': {
                'label': 'Switch to editor 4',
                'shortcut': '<Control>4',
                'callback': 'win.docview.switch_to_editor',
                'args': [4]
            },
            'switch_editor5': {
                'label': 'Switch to editor 5',
                'shortcut': '<Control>5',
                'callback': 'win.docview.switch_to_editor',
                'args': [5]
            },
            'switch_editor6': {
                'label': 'Switch to editor 6',
                'shortcut': '<Control>6',
                'callback': 'win.docview.switch_to_editor',
                'args': [6]
            },
            'switch_editor7': {
                'label': 'Switch to editor 7',
                'shortcut': '<Control>7',
                'callback': 'win.docview.switch_to_editor',
                'args': [7]
            },
            'switch_editor8': {
                'label': 'Switch to editor 8',
                'shortcut': '<Control>8',
                'callback': 'win.docview.switch_to_editor',
                'args': [8]
            },
            'switch_editor9': {
                'label': 'Switch to editor 9',
                'shortcut': '<Control>9',
                'callback': 'win.docview.switch_to_editor',
                'args': [9]
            }
        }
    },
    'editor': {
        'label': 'SQL Editor',
        'actions': {
            'run': {
                'label': 'Run SQL statement',
                'description': 'Run SQL statement at cursor',
                'icon': 'media-playback-start-symbolic',
                'shortcut': '<Control>Return',
                'callback': 'run_query'
            },
            'dbconnect': {
                'label': 'Connect',
                'description': 'Open or change database connection',
                'icon': 'gtk-connect',
                'shortcut': 'F9',
                'callback': 'assume_connection',
                'args': [True]
            },
            'dbdisconnect': {
                'label': 'Disconnect',
                'description': 'Close database connection',
                'icon': 'gtk-disconnect',
                'shortcut': 'F10',
                'callback': 'set_connection',
                'args': [None]
            },
            'openconnection': {
                'label': 'Open assigned connection',
                'description': 'Opens the currently assigned connection',
                'shortcut': 'F11',
                'callback': 'open_connection',
            },
            'format': {
                'label': 'Format SQL',
                'description': 'Format SQL statement at cursor',
                'icon': 'format-indent-more-symbolic',
                'shortcut': '<Alt>f',
                'callback': 'editor.format_statement'
            },
            'jump_next': {
                'label': 'Next statement',
                'description': 'Jump to next statement',
                'shortcut': '<Alt><Shift>Down',
                'callback': 'editor.jump_next'
            },
            'jump_prev': {
                'label': 'Previous statement',
                'description': 'Jump to previous statement',
                'shortcut': '<Alt><Shift>Up',
                'callback': 'editor.jump_prev'
            },
        }
    }
}
<commit_msg>Actually change shortcuts.<commit_after>commands = {
    'app': {
        'label': 'Application',
        'actions': {
            'neweditor': {
                'label': 'New SQL editor',
                'description': 'Open new SQL editor',
                'icon': 'document-new-symbolic',
                'shortcut': '<Control>N',
                'callback': 'win.docview.add_worksheet'
            },
            'switch_editor1': {
                'label': 'Switch to editor 1',
                'shortcut': '<Control>1',
                'callback': 'win.docview.switch_to_editor',
                'args': [1]
            },
            'switch_editor2': {
                'label': 'Switch to editor 2',
                'shortcut': '<Control>2',
                'callback': 'win.docview.switch_to_editor',
                'args': [2]
            },
            'switch_editor3': {
                'label': 'Switch to editor 3',
                'shortcut': '<Control>3',
                'callback': 'win.docview.switch_to_editor',
                'args': [3]
            },
            'switch_editor4': {
                'label': 'Switch to editor 4',
                'shortcut': '<Control>4',
                'callback': 'win.docview.switch_to_editor',
                'args': [4]
            },
            'switch_editor5': {
                'label': 'Switch to editor 5',
                'shortcut': '<Control>5',
                'callback': 'win.docview.switch_to_editor',
                'args': [5]
            },
            'switch_editor6': {
                'label': 'Switch to editor 6',
                'shortcut': '<Control>6',
                'callback': 'win.docview.switch_to_editor',
                'args': [6]
            },
            'switch_editor7': {
                'label': 'Switch to editor 7',
                'shortcut': '<Control>7',
                'callback': 'win.docview.switch_to_editor',
                'args': [7]
            },
            'switch_editor8': {
                'label': 'Switch to editor 8',
                'shortcut': '<Control>8',
                'callback': 'win.docview.switch_to_editor',
                'args': [8]
            },
            'switch_editor9': {
                'label': 'Switch to editor 9',
                'shortcut': '<Control>9',
                'callback': 'win.docview.switch_to_editor',
                'args': [9]
            }
        }
    },
    'editor': {
        'label': 'SQL Editor',
        'actions': {
            'run': {
                'label': 'Run SQL statement',
                'description': 'Run SQL statement at cursor',
                'icon': 'media-playback-start-symbolic',
                'shortcut': '<Control>Return',
                'callback': 'run_query'
            },
            'dbconnect': {
                'label': 'Connect',
                'description': 'Open or change database connection',
                'icon': 'gtk-connect',
                'shortcut': 'F9',
                'callback': 'assume_connection',
                'args': [True]
            },
            'openconnection': {
                'label': 'Open assigned connection',
                'description': 'Opens the currently assigned connection',
                'shortcut': 'F10',
                'callback': 'open_connection',
            },
            'dbdisconnect': {
                'label': 'Disconnect',
                'description': 'Close database connection',
                'icon': 'gtk-disconnect',
                'shortcut': 'F11',
                'callback': 'set_connection',
                'args': [None]
            },
            'format': {
                'label': 'Format SQL',
                'description': 'Format SQL statement at cursor',
                'icon': 'format-indent-more-symbolic',
                'shortcut': '<Alt>f',
                'callback': 'editor.format_statement'
            },
            'jump_next': {
                'label': 'Next statement',
                'description': 'Jump to next statement',
                'shortcut': '<Alt><Shift>Down',
                'callback': 'editor.jump_next'
            },
            'jump_prev': {
                'label': 'Previous statement',
                'description': 'Jump to previous statement',
                'shortcut': '<Alt><Shift>Up',
                'callback': 'editor.jump_prev'
            },
        }
    }
}
<|endoftext|>


<commit_before>from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

from contextlib import contextmanager
import ctypes.wintypes
import datetime
import errno
import logging
import os

try:
    import winreg
except ImportError:
    # py 2
    import _winreg as winreg

@contextmanager
def ignored(*exceptions):
    """Use contextlib to ignore all windows specific errors.
       These are generally encountered from missing registry keys,
       and can safely be ignored in most circumstances."""
    try:
        yield
    except exceptions:
        pass

fnf_exception = getattr(__builtins__,
                        'FileNotFoundError', WindowsError)

log = logging.getLogger(__name__)

# cyptes constants
CSIDL_PROFILE = 40
SHGFP_TYPE_CURRENT = 0

def _documents_folder():
    """ Get the users' documents folder, which is where R will place
        its default user-specific 'personal' library.

        Returns: full path of user library."""

    # first, check if the user has an R_USER variable initialized.
    documents_folder = _environ_path("R_USER")

    if not documents_folder:
        # Call SHGetFolderPath using ctypes.
        ctypes_buffer = ctypes.create_unicode_buffer(ctypes.wintypes.MAX_PATH)
        ctypes.windll.shell32.SHGetFolderPathW(
            0, CSIDL_PROFILE, 0, SHGFP_TYPE_CURRENT, ctypes_buffer)
        # This isn't a language-independent way, but CSIDL_PERSONAL gets
        # the wrong path.
        # TODO: Test in non-English locales.
        documents_folder = os.path.join(ctypes_buffer.value, "Documents")

    return documents_folder

def _environ_path(var=None):
    """ Check if an environment variable is an existing path."""
    path = None
    if var and var in os.environ:
        var_path = os.environ[var]
        if os.path.exists(var_path):
            path = var_path

    return path

def r_path():
    """Find R installation path from registry."""
    r_install_path = None

    # set an epoch for a Windows FILETIME object
    epoch = datetime.datetime(1601, 1, 1)

    root_keys = {
        'HKCU': winreg.HKEY_CURRENT_USER,
        'HKLM': winreg.HKEY_LOCAL_MACHINE
    }
    r_reg_paths = ["SOFTWARE\\R-core\\R",
                   "SOFTWARE\\R-core\\R64",
                   "SOFTWARE\\Wow6432Node\\R-Core\\R",
                   "SOFTWARE\\Wow6432Node\\R-Core\\R64"]

    for (key_name, root_key) in list(root_keys.items()):
        for r_path in r_reg_paths:
            r_reg = None

            try:
                log.info("OpenKey on {}, with READ + WOW64\n".format(
                    r_path))
                r_reg = winreg.OpenKey(root_key, r_path,
                                       0, (winreg.KEY_WOW64_64KEY +
                                           winreg.KEY_READ))
            except fnf_exception as error:
                log.debug("Exception generated: {}".format(error))
                if error.errno == errno.ENOENT:
                    pass
                else:
                    raise

            if r_reg:
                log.info("Successfully found {}".format(r_path))

                try:
                    log.info("Looking for InstallPath.")
                    r_install_path = winreg.QueryValueEx(r_reg, "InstallPath")[0]
                except fnf_exception as error:
                    log.debug("Exception generated: {}".format(error))
                    if error.errno == errno.ENOENT:
                        pass
                    else:
                        raise

                if not r_install_path:
                    log.debug("Top-level install path not defined. " +
                              "Checking version-specific locations.")
                    # Can't find the install path as a top-level value.
                    # Inspect the children keys for versions, and use the most
                    # recently installed one as the correct R installation.
                    max_time = epoch

                    for pos in range(10):
                        # TODO ensure this is robust to errors
                        with ignored(WindowsError):
                            log.info("checking EnumKey pos {}".format(pos))
                            r_base_key = winreg.EnumKey(r_reg, pos)

                            # test for the right path based on age
                            if r_base_key:
                                r_version_key = "{}\\{}".format(
                                    r_path, r_base_key)
                                r_version_reg = winreg.OpenKey(
                                    root_key, r_version_key, 0,
                                    (winreg.KEY_READ |
                                     winreg.KEY_WOW64_64KEY))
                                r_install_path = winreg.QueryValueEx(
                                    r_version_reg, "InstallPath")[0]
                                r_version_info = winreg.QueryInfoKey(
                                    r_version_reg)
                                r_install_time = epoch + datetime.timedelta(
                                    microseconds=r_version_info[2]/10)
                                if max_time < r_install_time:
                                    max_time = r_install_time
    log.info("Final R install path: {}".format(r_install_path))
    return r_install_path

r_install_path = r_path()

def r_version():
    r_version = None
    r_path_l = r_install_path
    if r_path_l is not None:
        r_version = r_path_l.split('-')[1]
    return r_version

r_version_info = r_version()

def r_all_lib_paths():
    """ Package library, locates all known library
        paths used for R packages."""

    libs_path = []
    # check R_LIBS_USER first
    if _environ_path("R_LIBS_USER"):
        libs_path.append(_environ_path("R_LIBS_USER"))

    # user's R library in Documents/R/win-library/R-x.x/
    (r_major, r_minor, r_patch) = r_version_info.split(".")

    r_user_library_path = os.path.join(
        _documents_folder(), "R", "win-library",
        "{}.{}".format(r_major, r_minor))
    if os.path.exists(r_user_library_path):
        libs_path.append(r_user_library_path)

    # Next, check the value of R_LIBS -- users may set this
    # instead of the (more specific) R_LIBS_USER
    if _environ_path("R_LIBS"):
        libs_path.append(_environ_path("R_LIBS"))

    # R library in Program Files/R-x.xx/library
    # NOTE: Requires elevated privileges to write to
    if r_install_path is not None:
        r_install_lib_path = os.path.join(
            r_install_path, "library")

        if os.path.exists(r_install_lib_path):
            libs_path.append(r_install_lib_path)

    return libs_path

r_all_library_paths = r_all_lib_paths()

def r_lib_path():
    """ Package library, locates the highest-priority
        library path used for R packages."""
    lib_path = None
    all_libs = r_all_library_paths
    if len(all_libs) > 0:
        lib_path = all_libs[0]
    return lib_path

r_library_path = r_lib_path()

def r_pkg_path():
    """
    Package path search. Locations searched:
     - HKCU\\Software\\Esri\\ArcGISPro\\RintegrationProPackagePath
     - [MYDOCUMENTS]/R/win-library/[3-9].[0-9]/ - default for user R packages
     - [ArcGIS]/Resources/Rintegration/arcgisbinding
    """
    # NOTE to be 100% robust, this may be better implemented as a call to
    #      R, and parsing out the .libPaths() results.
    package_path = None
    package_name = 'arcgisbinding'

    root_key = winreg.HKEY_CURRENT_USER
    reg_path = "SOFTWARE\\Esri\\ArcGISPro"
    package_key = 'RintegrationProPackagePath'
    pro_reg = None

    try:
        # find the key, 64- or 32-bit we want it all
        pro_reg = winreg.OpenKey(
            root_key, reg_path, 0,
            (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
    except fnf_exception as error:
        if error.errno == errno.ENOENT:
            pass
        else:
            raise

    if pro_reg:
        try:
            # returns a tuple of (value, type)
            package_path_key = winreg.QueryValueEx(pro_reg, package_key)
            package_path_raw = package_path_key[0]
            if os.path.exists(package_path_raw):
                package_path = package_path_raw
        except fnf_exception as error:
            if error.errno == errno.ENOENT:
                pass
            else:
                raise

    # iterate over all known library path locations,
    # and check for our package in each.
    for lib_path in r_all_library_paths:
        possible_package_path = os.path.join(lib_path, package_name)
        if os.path.exists(possible_package_path):
            package_path = possible_package_path
            # we want the highest-priority library, stop here
            break

    # fallback -- <ArcGIS Install>/Rintegration/arcgisbinding
    if not package_path:
        import arcpy
        arc_install_dir = arcpy.GetInstallInfo()['InstallDir']
        arc_package_dir = os.path.join(
            arc_install_dir, 'Rintegration', package_name)
        if os.path.exists(arc_package_dir):
            package_path = arc_package_dir

    return package_path

r_package_path = r_pkg_path()

def r_pkg_version():
    version = None
    r_package_path = r_pkg_path()
    if r_package_path:
        desc_path = os.path.join(r_package_path, 'DESCRIPTION')
        if os.path.exists(desc_path):
            with open(desc_path) as desc_f:
                for line in desc_f:
                    try:
                        (key, value_raw) = line.strip().split(':')
                    except:
                        # gulp
                        pass
                    if key == 'Version':
                        version = value_raw.strip()
    return version

r_package_version = r_pkg_version()

def arcmap_exists(version=None):
    root_key = winreg.HKEY_CURRENT_USER
    reg_base = "SOFTWARE\\Esri"
    if not version:
        version = "10.3"
    package_key = "Desktop{}".format(version)

    reg_path = "{}\\{}".format(reg_base, package_key)
    arcmap_reg = None
    installed = False
    try:
        # find the key, 64- or 32-bit we want it all
        arcmap_reg = winreg.OpenKey(
            root_key, reg_path, 0,
            (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
    except fnf_exception as error:
        if error.errno == errno.ENOENT:
            pass
        else:
            raise

    if arcmap_reg:
        installed = True

    return installed

def arcmap_path(version=None):
    """Path to ArcGIS Installation."""

    arcmap_path = None

    if not version:
        version = "10.3"
    package_key = "Desktop{}".format(version)

    root_key = winreg.HKEY_LOCAL_MACHINE
    arc_reg_paths = [
        "SOFTWARE\\ESRI\\{}".format(package_key),
        "SOFTWARE\\Wow6432Node\\ESRI\\{}".format(package_key)
    ]

    arcmap_reg = None
    for reg_path in arc_reg_paths:
        try:
            # find the key, 64- or 32-bit we want it all
            arcmap_reg = winreg.OpenKey(
                root_key, reg_path, 0,
                (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
        except fnf_exception as error:
            if error.errno == errno.ENOENT:
                pass
            else:
                raise

        if arcmap_reg:
            try:
                # returns a tuple of (value, type)
                arcmap_path_key = winreg.QueryValueEx(arcmap_reg, "InstallDir")
                arcmap_path_raw = arcmap_path_key[0]
                if os.path.exists(arcmap_path_raw):
                    arcmap_path = arcmap_path_raw.strip('\\')
            except fnf_exception as error:
                if error.errno == errno.ENOENT:
                    pass
                else:
                    raise

    return arcmap_path

arcmap_install_path = arcmap_path()
<commit_msg>Also search R_LIBS_SITE<commit_after>from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

from contextlib import contextmanager
import ctypes.wintypes
import datetime
import errno
import logging
import os

try:
    import winreg
except ImportError:
    # py 2
    import _winreg as winreg

@contextmanager
def ignored(*exceptions):
    """Use contextlib to ignore all windows specific errors.
       These are generally encountered from missing registry keys,
       and can safely be ignored in most circumstances."""
    try:
        yield
    except exceptions:
        pass

fnf_exception = getattr(__builtins__,
                        'FileNotFoundError', WindowsError)

log = logging.getLogger(__name__)

# cyptes constants
CSIDL_PROFILE = 40
SHGFP_TYPE_CURRENT = 0

def _documents_folder():
    """ Get the users' documents folder, which is where R will place
        its default user-specific 'personal' library.

        Returns: full path of user library."""

    # first, check if the user has an R_USER variable initialized.
    documents_folder = _environ_path("R_USER")

    if not documents_folder:
        # Call SHGetFolderPath using ctypes.
        ctypes_buffer = ctypes.create_unicode_buffer(ctypes.wintypes.MAX_PATH)
        ctypes.windll.shell32.SHGetFolderPathW(
            0, CSIDL_PROFILE, 0, SHGFP_TYPE_CURRENT, ctypes_buffer)
        # This isn't a language-independent way, but CSIDL_PERSONAL gets
        # the wrong path.
        # TODO: Test in non-English locales.
        documents_folder = os.path.join(ctypes_buffer.value, "Documents")

    return documents_folder

def _environ_path(var=None):
    """ Check if an environment variable is an existing path."""
    path = None
    if var and var in os.environ:
        var_path = os.environ[var]
        if os.path.exists(var_path):
            path = var_path

    return path

def r_path():
    """Find R installation path from registry."""
    r_install_path = None

    # set an epoch for a Windows FILETIME object
    epoch = datetime.datetime(1601, 1, 1)

    root_keys = {
        'HKCU': winreg.HKEY_CURRENT_USER,
        'HKLM': winreg.HKEY_LOCAL_MACHINE
    }
    r_reg_paths = ["SOFTWARE\\R-core\\R",
                   "SOFTWARE\\R-core\\R64",
                   "SOFTWARE\\Wow6432Node\\R-Core\\R",
                   "SOFTWARE\\Wow6432Node\\R-Core\\R64"]

    for (key_name, root_key) in list(root_keys.items()):
        for r_path in r_reg_paths:
            r_reg = None

            try:
                log.info("OpenKey on {}, with READ + WOW64\n".format(
                    r_path))
                r_reg = winreg.OpenKey(root_key, r_path,
                                       0, (winreg.KEY_WOW64_64KEY +
                                           winreg.KEY_READ))
            except fnf_exception as error:
                log.debug("Exception generated: {}".format(error))
                if error.errno == errno.ENOENT:
                    pass
                else:
                    raise

            if r_reg:
                log.info("Successfully found {}".format(r_path))

                try:
                    log.info("Looking for InstallPath.")
                    r_install_path = winreg.QueryValueEx(r_reg, "InstallPath")[0]
                except fnf_exception as error:
                    log.debug("Exception generated: {}".format(error))
                    if error.errno == errno.ENOENT:
                        pass
                    else:
                        raise

                if not r_install_path:
                    log.debug("Top-level install path not defined. " +
                              "Checking version-specific locations.")
                    # Can't find the install path as a top-level value.
                    # Inspect the children keys for versions, and use the most
                    # recently installed one as the correct R installation.
                    max_time = epoch

                    for pos in range(10):
                        # TODO ensure this is robust to errors
                        with ignored(WindowsError):
                            log.info("checking EnumKey pos {}".format(pos))
                            r_base_key = winreg.EnumKey(r_reg, pos)

                            # test for the right path based on age
                            if r_base_key:
                                r_version_key = "{}\\{}".format(
                                    r_path, r_base_key)
                                r_version_reg = winreg.OpenKey(
                                    root_key, r_version_key, 0,
                                    (winreg.KEY_READ |
                                     winreg.KEY_WOW64_64KEY))
                                r_install_path = winreg.QueryValueEx(
                                    r_version_reg, "InstallPath")[0]
                                r_version_info = winreg.QueryInfoKey(
                                    r_version_reg)
                                r_install_time = epoch + datetime.timedelta(
                                    microseconds=r_version_info[2]/10)
                                if max_time < r_install_time:
                                    max_time = r_install_time
    log.info("Final R install path: {}".format(r_install_path))
    return r_install_path

r_install_path = r_path()

def r_version():
    r_version = None
    r_path_l = r_install_path
    if r_path_l is not None:
        r_version = r_path_l.split('-')[1]
    return r_version

r_version_info = r_version()

def r_all_lib_paths():
    """ Package library, locates all known library
        paths used for R packages."""

    libs_path = []
    # check R_LIBS_USER first
    if _environ_path("R_LIBS_USER"):
        libs_path.append(_environ_path("R_LIBS_USER"))

    # user's R library in Documents/R/win-library/R-x.x/
    (r_major, r_minor, r_patch) = r_version_info.split(".")

    r_user_library_path = os.path.join(
        _documents_folder(), "R", "win-library",
        "{}.{}".format(r_major, r_minor))
    if os.path.exists(r_user_library_path):
        libs_path.append(r_user_library_path)

    # Next, check the value of R_LIBS -- users may set this
    # instead of the (more specific) R_LIBS_USER
    if _environ_path("R_LIBS"):
        libs_path.append(_environ_path("R_LIBS"))

    # lastly, check for possible site libraries.
    # NOTE: Requires elevated privileges to write to

    if _environ_path("R_LIBS_SITE"):
        libs_path.append(_environ_path("R_LIBS_SITE"))

    # R library in Program Files/R-x.xx/library
    if r_install_path is not None:
        r_install_lib_path = os.path.join(
            r_install_path, "library")

        if os.path.exists(r_install_lib_path):
            libs_path.append(r_install_lib_path)

    return libs_path

r_all_library_paths = r_all_lib_paths()

def r_lib_path():
    """ Package library, locates the highest-priority
        library path used for R packages."""
    lib_path = None
    all_libs = r_all_library_paths
    if len(all_libs) > 0:
        lib_path = all_libs[0]
    return lib_path

r_library_path = r_lib_path()

def r_pkg_path():
    """
    Package path search. Locations searched:
     - HKCU\\Software\\Esri\\ArcGISPro\\RintegrationProPackagePath
     - [MYDOCUMENTS]/R/win-library/[3-9].[0-9]/ - default for user R packages
     - [ArcGIS]/Resources/Rintegration/arcgisbinding
    """
    package_path = None
    package_name = 'arcgisbinding'

    root_key = winreg.HKEY_CURRENT_USER
    reg_path = "SOFTWARE\\Esri\\ArcGISPro"
    package_key = 'RintegrationProPackagePath'
    pro_reg = None

    try:
        # find the key, 64- or 32-bit we want it all
        pro_reg = winreg.OpenKey(
            root_key, reg_path, 0,
            (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
    except fnf_exception as error:
        if error.errno == errno.ENOENT:
            pass
        else:
            raise

    if pro_reg:
        try:
            # returns a tuple of (value, type)
            package_path_key = winreg.QueryValueEx(pro_reg, package_key)
            package_path_raw = package_path_key[0]
            if os.path.exists(package_path_raw):
                package_path = package_path_raw
        except fnf_exception as error:
            if error.errno == errno.ENOENT:
                pass
            else:
                raise

    # iterate over all known library path locations,
    # and check for our package in each.
    for lib_path in r_all_library_paths:
        possible_package_path = os.path.join(lib_path, package_name)
        if os.path.exists(possible_package_path):
            package_path = possible_package_path
            # we want the highest-priority library, stop here
            break

    # fallback -- <ArcGIS Install>/Rintegration/arcgisbinding
    if not package_path:
        import arcpy
        arc_install_dir = arcpy.GetInstallInfo()['InstallDir']
        arc_package_dir = os.path.join(
            arc_install_dir, 'Rintegration', package_name)
        if os.path.exists(arc_package_dir):
            package_path = arc_package_dir

    return package_path

r_package_path = r_pkg_path()

def r_pkg_version():
    version = None
    r_package_path = r_pkg_path()
    if r_package_path:
        desc_path = os.path.join(r_package_path, 'DESCRIPTION')
        if os.path.exists(desc_path):
            with open(desc_path) as desc_f:
                for line in desc_f:
                    try:
                        (key, value_raw) = line.strip().split(':')
                    except:
                        # gulp
                        pass
                    if key == 'Version':
                        version = value_raw.strip()
    return version

r_package_version = r_pkg_version()

def arcmap_exists(version=None):
    root_key = winreg.HKEY_CURRENT_USER
    reg_base = "SOFTWARE\\Esri"
    if not version:
        version = "10.3"
    package_key = "Desktop{}".format(version)

    reg_path = "{}\\{}".format(reg_base, package_key)
    arcmap_reg = None
    installed = False
    try:
        # find the key, 64- or 32-bit we want it all
        arcmap_reg = winreg.OpenKey(
            root_key, reg_path, 0,
            (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
    except fnf_exception as error:
        if error.errno == errno.ENOENT:
            pass
        else:
            raise

    if arcmap_reg:
        installed = True

    return installed

def arcmap_path(version=None):
    """Path to ArcGIS Installation."""

    arcmap_path = None

    if not version:
        version = "10.3"
    package_key = "Desktop{}".format(version)

    root_key = winreg.HKEY_LOCAL_MACHINE
    arc_reg_paths = [
        "SOFTWARE\\ESRI\\{}".format(package_key),
        "SOFTWARE\\Wow6432Node\\ESRI\\{}".format(package_key)
    ]

    arcmap_reg = None
    for reg_path in arc_reg_paths:
        try:
            # find the key, 64- or 32-bit we want it all
            arcmap_reg = winreg.OpenKey(
                root_key, reg_path, 0,
                (winreg.KEY_READ | winreg.KEY_WOW64_64KEY))
        except fnf_exception as error:
            if error.errno == errno.ENOENT:
                pass
            else:
                raise

        if arcmap_reg:
            try:
                # returns a tuple of (value, type)
                arcmap_path_key = winreg.QueryValueEx(arcmap_reg, "InstallDir")
                arcmap_path_raw = arcmap_path_key[0]
                if os.path.exists(arcmap_path_raw):
                    arcmap_path = arcmap_path_raw.strip('\\')
            except fnf_exception as error:
                if error.errno == errno.ENOENT:
                    pass
                else:
                    raise

    return arcmap_path

arcmap_install_path = arcmap_path()
<|endoftext|>


<commit_before>import numpy as np

class Element:
    """Finite element."""

    maxdeg=0 # maximum polynomial degree; for determining quadrature
    dim=0 # spatial dimension
    tdim=0 # target dimension
    torder=0 # target tensorial order

    """
    The discretized field is a mapping
       U : C^dim -> C^(tdim x ... x tdim)
    where the product is taken 'torder' times.
    """

    # number of ...
    n_dofs=0 # nodal dofs
    i_dofs=0 # interior dofs
    f_dofs=0 # facet dofs (2d and 3d only)
    e_dofs=0 # edge dofs (3d only)

    def lbasis(self,X,i):
        """Returns local basis functions evaluated at some local points."""
        raise NotImplementedError("Element.lbasis: local basis (lbasis) not implemented!")

    def gbasis(self,X,i,tind):
        """Returns global basis functions evaluated at some local points."""
        raise NotImplementedError("Element.gbasis: local basis (lbasis) not implemented!")

class ElementH1(Element):
    """H1 conforming finite element."""

    def gbasis(self,mapping,X,i,tind):
        if isinstance(X,dict):
            [phi,dphi]=self.lbasis(X,i)
            u=phi
            du={}
            invDF=mapping.invDF(X,tind)

            if mapping.dim==2:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]
            elif mapping.dim==3:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]+invDF[2][0]*dphi[2]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]+invDF[2][1]*dphi[2]
                du[2]=invDF[0][2]*dphi[0]+invDF[1][2]*dphi[1]+invDF[2][2]*dphi[2]
            else:
                raise NotImplementedError("ElementH1.gbasis: not implemented for the given dim.")
            ddu=None # TODO fix ddu (for H1 element, Laplacian?)

        else:
            x={}
            x[0]=X[0,:]
            if mapping.dim>=2:
                x[1]=X[1,:]
            if mapping.dim>=3:
                x[2]=X[2,:]
            [phi,dphi]=self.lbasis(x,i)
            u=np.tile(phi,(len(tind),1))
            du={}
            invDF=mapping.invDF(X,tind)

            if mapping.dim==2:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]
            elif mapping.dim==3:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]+invDF[2][0]*dphi[2]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]+invDF[2][1]*dphi[2]
                du[2]=invDF[0][2]*dphi[0]+invDF[1][2]*dphi[1]+invDF[2][2]*dphi[2]
            else:
                raise NotImplementedError("ElementH1.gbasis: not implemented for the given dim.")
            ddu=None # TODO fix ddu (for H1 element, Laplacian?)

        return u,du,ddu

class ElementTriPpFast(ElementH1):
    def __init__(self,p):
        self.p=p
        self.maxdeg=p
        self.n_dofs=1
        self.f_dofs=np.max([p-1,0])
        self.i_dofs=np.max([(p-1)*(p-2)/2,0])

        self.nbdofs=3*self.n_dofs+3*self.f_dofs+self.i_dofs
        # TODO build bfuns using e.g. sympy and make handles of those.

class ElementTriPp(ElementH1):
    """A somewhat slow implementation of p-refinements for triangular mesh."""
    def __init__(self,p):
        self.p=p
        self.maxdeg=p
        self.n_dofs=1
        self.f_dofs=np.max([p-1,0])
        self.i_dofs=np.max([(p-1)*(p-2)/2,0])

        self.nbdofs=3*self.n_dofs+3*self.f_dofs+self.i_dofs

    def intlegpoly(self,x,n):
        """Generate integrated Legendre polynomials."""
        n=n+1

        P={}
        P[0]=np.ones(x.shape)
        if n>1:
            P[1]=x

        for i in np.arange(1,n):
            P[i+1]=((2.*i+1.)/(i+1.))*x*P[i]-(i/(i+1.))*P[i-1]

        iP={}
        iP[0]=np.ones(x.shape)
        if n>1:
            iP[1]=x

        for i in np.arange(1,n-1):
            iP[i+1]=(P[i+1]-P[i-1])/(2.*i+1.)

        dP={}
        dP[0]=np.zeros(x.shape)
        for i in np.arange(1,n):
            dP[i]=P[i-1]

        return iP,dP

    def lbasis(self,X,n):
        """Evaluate n'th Lagrange basis of order self.p."""
        p=self.p

        if len(X)!=2:
            raise NotImplementedError("ElementTriPp: not implemented for the given dimension of X.")

        phi={}
        phi[0]=1.-X[0]-X[1]
        phi[1]=X[0]
        phi[2]=X[1]

        # local basis function gradients TODO fix these somehow
        gradphi_x={}
        gradphi_x[0]=-1.*np.ones(X[0].shape)
        gradphi_x[1]=1.*np.ones(X[0].shape)
        gradphi_x[2]=np.zeros(X[0].shape)

        gradphi_y={}
        gradphi_y[0]=-1.*np.ones(X[0].shape)
        gradphi_y[1]=np.zeros(X[0].shape)
        gradphi_y[2]=1.*np.ones(X[0].shape)

        if n<=2:
            # return first three
            dphi={}
            dphi[0]=gradphi_x[n]
            dphi[1]=gradphi_y[n]
            return phi[n],dphi

        # use same ordering as in mesh
        e=np.array([[0,1],[1,2],[0,2]]).T
        offset=3

        # define edge basis functions
        if(p>1):
            for i in range(3):
                eta=phi[e[1,i]]-phi[e[0,i]]
                deta_x=gradphi_x[e[1,i]]-gradphi_x[e[0,i]]
                deta_y=gradphi_y[e[1,i]]-gradphi_y[e[0,i]]

                # generate integrated Legendre polynomials
                [P,dP]=self.intlegpoly(eta,p-2)
                   #print len(P)
                   #print len(dP)

                for j in range(len(P)):
                    phi[offset]=phi[e[0,i]]*phi[e[1,i]]*P[j]
                    gradphi_x[offset]=gradphi_x[e[0,i]]*phi[e[1,i]]*P[j]+\
                                      gradphi_x[e[1,i]]*phi[e[0,i]]*P[j]+\
                                      deta_x*phi[e[0,i]]*phi[e[1,i]]*dP[j]
                    gradphi_y[offset]=gradphi_y[e[0,i]]*phi[e[1,i]]*P[j]+\
                                      gradphi_y[e[1,i]]*phi[e[0,i]]*P[j]+\
                                      deta_y*phi[e[0,i]]*phi[e[1,i]]*dP[j]
                    if offset==n:
                        # return if computed
                        dphi={}
                        dphi[0]=gradphi_x[n]
                        dphi[1]=gradphi_y[n]
                        return phi[n],dphi
                    offset=offset+1

        # define interior basis functions
        if(p>2):
            B={}
            dB_x={}
            dB_y={}
            if(p>3):
                pm3=ElementTriPp(p-3)
                for itr in range(pm3.nbdofs):
                    pphi,pdphi=self.lbasis(X,itr)
                    B[itr]=pphi
                    dB_x[itr]=pdphi[0]
                    dB_y[itr]=pdphi[1]
                N=pm3.nbdofs
            else:
                B[0]=np.ones(X[0].shape)
                dB_x[0]=np.zeros(X[0].shape)
                dB_y[0]=np.zeros(X[0].shape)
                N=1

            bubble=phi[0]*phi[1]*phi[2]
            dbubble_x=gradphi_x[0]*phi[1]*phi[2]+\
                      gradphi_x[1]*phi[2]*phi[0]+\
                      gradphi_x[2]*phi[0]*phi[1]
            dbubble_y=gradphi_y[0]*phi[1]*phi[2]+\
                      gradphi_y[1]*phi[2]*phi[0]+\
                      gradphi_y[2]*phi[0]*phi[1]

            for i in range(N):
                phi[offset]=bubble*B[i]
                gradphi_x[offset]=dbubble_x*B[i]+dB_x[i]*bubble
                gradphi_y[offset]=dbubble_y*B[i]+dB_y[i]*bubble
                if offset==n:
                    # return if computed
                    dphi={}
                    dphi[0]=gradphi_x[n]
                    dphi[1]=gradphi_y[n]
                    return phi[n],dphi
                offset=offset+1

        raise IndexError("ElementTriPp.lbasis: reached end of lbasis without returning anything.")

class ElementP1(ElementH1):

    n_dofs=1
    maxdeg=1

    def __init__(self,dim=1):
        self.dim=dim

    def lbasis(self,X,i):
        # TODO implement for 1D

        if self.dim==2:
            phi={
                0:lambda x,y: 1-x-y,
                1:lambda x,y: x,
                2:lambda x,y: y
                }[i](X[0],X[1])

            dphi={}
            dphi[0]={
                    0:lambda x,y: -1+0*x,
                    1:lambda x,y: 1+0*x,
                    2:lambda x,y: 0*x
                    }[i](X[0],X[1])
            dphi[1]={
                    0:lambda x,y: -1+0*x,
                    1:lambda x,y: 0*x,
                    2:lambda x,y: 1+0*x
                    }[i](X[0],X[1])
        elif self.dim==3:
            phi={
                0:lambda x,y,z: 1-x-y-z,
                1:lambda x,y,z: x,
                2:lambda x,y,z: y,
                3:lambda x,y,z: z,
                }[i](X[0],X[1],X[2])

            dphi={}
            dphi[0]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 1+0*x,
                    2:lambda x,y,z: 0*x,
                    3:lambda x,y,z: 0*x
                    }[i](X[0],X[1],X[2])
            dphi[1]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 0*x,
                    2:lambda x,y,z: 1+0*x,
                    3:lambda x,y,z: 0*x
                    }[i](X[0],X[1],X[2])
            dphi[2]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 0*x,
                    2:lambda x,y,z: 0*x,
                    3:lambda x,y,z: 1+0*x
                    }[i](X[0],X[1],X[2])
        else:
            raise NotImplementedError("ElementP1.lbasis: not implemented for the given X.shape[0].")

        return phi,dphi

<commit_msg>Started on Lagrange Pp<commit_after>import numpy as np

class Element:
    """Finite element."""

    maxdeg=0 # maximum polynomial degree; for determining quadrature
    dim=0 # spatial dimension
    tdim=0 # target dimension
    torder=0 # target tensorial order

    """
    The discretized field is a mapping
       U : C^dim -> C^(tdim x ... x tdim)
    where the product is taken 'torder' times.
    """

    # number of ...
    n_dofs=0 # nodal dofs
    i_dofs=0 # interior dofs
    f_dofs=0 # facet dofs (2d and 3d only)
    e_dofs=0 # edge dofs (3d only)

    def lbasis(self,X,i):
        """Returns local basis functions evaluated at some local points."""
        raise NotImplementedError("Element.lbasis: local basis (lbasis) not implemented!")

    def gbasis(self,X,i,tind):
        """Returns global basis functions evaluated at some local points."""
        raise NotImplementedError("Element.gbasis: local basis (lbasis) not implemented!")

class ElementH1(Element):
    """H1 conforming finite element."""

    def gbasis(self,mapping,X,i,tind):
        if isinstance(X,dict):
            [phi,dphi]=self.lbasis(X,i)
            u=phi
            du={}
            invDF=mapping.invDF(X,tind)

            if mapping.dim==2:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]
            elif mapping.dim==3:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]+invDF[2][0]*dphi[2]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]+invDF[2][1]*dphi[2]
                du[2]=invDF[0][2]*dphi[0]+invDF[1][2]*dphi[1]+invDF[2][2]*dphi[2]
            else:
                raise NotImplementedError("ElementH1.gbasis: not implemented for the given dim.")
            ddu=None # TODO fix ddu (for H1 element, Laplacian?)

        else:
            x={}
            x[0]=X[0,:]
            if mapping.dim>=2:
                x[1]=X[1,:]
            if mapping.dim>=3:
                x[2]=X[2,:]
            [phi,dphi]=self.lbasis(x,i)
            u=np.tile(phi,(len(tind),1))
            du={}
            invDF=mapping.invDF(X,tind)

            if mapping.dim==2:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]
            elif mapping.dim==3:
                du[0]=invDF[0][0]*dphi[0]+invDF[1][0]*dphi[1]+invDF[2][0]*dphi[2]
                du[1]=invDF[0][1]*dphi[0]+invDF[1][1]*dphi[1]+invDF[2][1]*dphi[2]
                du[2]=invDF[0][2]*dphi[0]+invDF[1][2]*dphi[1]+invDF[2][2]*dphi[2]
            else:
                raise NotImplementedError("ElementH1.gbasis: not implemented for the given dim.")
            ddu=None # TODO fix ddu (for H1 element, Laplacian?)

        return u,du,ddu

class ElementTriPpFast(ElementH1):
    def __init__(self,p):
        self.p=p
        self.maxdeg=p
        self.n_dofs=1
        self.f_dofs=np.max([p-1,0])
        self.i_dofs=np.max([(p-1)*(p-2)/2,0])

        self.nbdofs=3*self.n_dofs+3*self.f_dofs+self.i_dofs
        # TODO build bfuns using e.g. sympy and make handles of those.

class ElementTriPp(ElementH1):
    """A somewhat slow implementation of p-refinements for triangular mesh."""
    def __init__(self,p):
        self.p=p
        self.maxdeg=p
        self.n_dofs=1
        self.f_dofs=np.max([p-1,0])
        self.i_dofs=np.max([(p-1)*(p-2)/2,0])

        self.nbdofs=3*self.n_dofs+3*self.f_dofs+self.i_dofs

    def intlegpoly(self,x,n):
        """Generate integrated Legendre polynomials."""
        n=n+1

        P={}
        P[0]=np.ones(x.shape)
        if n>1:
            P[1]=x

        for i in np.arange(1,n):
            P[i+1]=((2.*i+1.)/(i+1.))*x*P[i]-(i/(i+1.))*P[i-1]

        iP={}
        iP[0]=np.ones(x.shape)
        if n>1:
            iP[1]=x

        for i in np.arange(1,n-1):
            iP[i+1]=(P[i+1]-P[i-1])/(2.*i+1.)

        dP={}
        dP[0]=np.zeros(x.shape)
        for i in np.arange(1,n):
            dP[i]=P[i-1]

        return iP,dP

    def lbasis(self,X,n):
        """Evaluate n'th Lagrange basis of order self.p."""
        p=self.p

        if len(X)!=2:
            raise NotImplementedError("ElementTriPp: not implemented for the given dimension of X.")

        phi={}
        phi[0]=1.-X[0]-X[1]
        phi[1]=X[0]
        phi[2]=X[1]

        # local basis function gradients TODO fix these somehow
        gradphi_x={}
        gradphi_x[0]=-1.*np.ones(X[0].shape)
        gradphi_x[1]=1.*np.ones(X[0].shape)
        gradphi_x[2]=np.zeros(X[0].shape)

        gradphi_y={}
        gradphi_y[0]=-1.*np.ones(X[0].shape)
        gradphi_y[1]=np.zeros(X[0].shape)
        gradphi_y[2]=1.*np.ones(X[0].shape)

        if n<=2:
            # return first three
            dphi={}
            dphi[0]=gradphi_x[n]
            dphi[1]=gradphi_y[n]
            return phi[n],dphi

        # use same ordering as in mesh
        e=np.array([[0,1],[1,2],[0,2]]).T
        offset=3

        # define edge basis functions
        if(p>1):
            for i in range(3):
                eta=phi[e[1,i]]-phi[e[0,i]]
                deta_x=gradphi_x[e[1,i]]-gradphi_x[e[0,i]]
                deta_y=gradphi_y[e[1,i]]-gradphi_y[e[0,i]]

                # generate integrated Legendre polynomials
                [P,dP]=self.intlegpoly(eta,p-2)
                   #print len(P)
                   #print len(dP)

                for j in range(len(P)):
                    phi[offset]=phi[e[0,i]]*phi[e[1,i]]*P[j]
                    gradphi_x[offset]=gradphi_x[e[0,i]]*phi[e[1,i]]*P[j]+\
                                      gradphi_x[e[1,i]]*phi[e[0,i]]*P[j]+\
                                      deta_x*phi[e[0,i]]*phi[e[1,i]]*dP[j]
                    gradphi_y[offset]=gradphi_y[e[0,i]]*phi[e[1,i]]*P[j]+\
                                      gradphi_y[e[1,i]]*phi[e[0,i]]*P[j]+\
                                      deta_y*phi[e[0,i]]*phi[e[1,i]]*dP[j]
                    if offset==n:
                        # return if computed
                        dphi={}
                        dphi[0]=gradphi_x[n]
                        dphi[1]=gradphi_y[n]
                        return phi[n],dphi
                    offset=offset+1

        # define interior basis functions
        if(p>2):
            B={}
            dB_x={}
            dB_y={}
            if(p>3):
                pm3=ElementTriPp(p-3)
                for itr in range(pm3.nbdofs):
                    pphi,pdphi=self.lbasis(X,itr)
                    B[itr]=pphi
                    dB_x[itr]=pdphi[0]
                    dB_y[itr]=pdphi[1]
                N=pm3.nbdofs
            else:
                B[0]=np.ones(X[0].shape)
                dB_x[0]=np.zeros(X[0].shape)
                dB_y[0]=np.zeros(X[0].shape)
                N=1

            bubble=phi[0]*phi[1]*phi[2]
            dbubble_x=gradphi_x[0]*phi[1]*phi[2]+\
                      gradphi_x[1]*phi[2]*phi[0]+\
                      gradphi_x[2]*phi[0]*phi[1]
            dbubble_y=gradphi_y[0]*phi[1]*phi[2]+\
                      gradphi_y[1]*phi[2]*phi[0]+\
                      gradphi_y[2]*phi[0]*phi[1]

            for i in range(N):
                phi[offset]=bubble*B[i]
                gradphi_x[offset]=dbubble_x*B[i]+dB_x[i]*bubble
                gradphi_y[offset]=dbubble_y*B[i]+dB_y[i]*bubble
                if offset==n:
                    # return if computed
                    dphi={}
                    dphi[0]=gradphi_x[n]
                    dphi[1]=gradphi_y[n]
                    return phi[n],dphi
                offset=offset+1

        raise IndexError("ElementTriPp.lbasis: reached end of lbasis without returning anything.")

class ElementTriLagrangePp(ElementH1):

    def __init__(self,p=1):
        self.p=p
        self.maxdeg=p

        ndofs=int((p+1)*(p+2)/2.)

        V=np.zeros((ndofs,ndofs))

        nx=np.zeros(ndofs)
        ny=np.zeros(ndofs)

        ps={}

        ktr=0
        ltr=0
        for itr in np.linspace(0.0,1.0,p+1):
            for jtr in np.linspace(0.0,1.0-itr,p+1-ktr):
                nx[ltr]=itr
                ny[ltr]=jtr
                ltr=ltr+1
            ktr=ktr+1

        print nx
        print ny

        ktr=0
        for itr in range(p+1):
            for jtr in range(p+1):
                if itr+jtr<=p:
                    poly=lambda x,y: (x**itr)*(y**jtr)
                    V[ktr,:]=poly(nx,ny)
                    ps[ktr]=(itr,jtr)
                    print ps[ktr]
                    ktr=ktr+1

        c=np.zeros((ndofs,ndofs))
        V=np.linalg.inv(V)

        bfuns={}

        for itr in range(ndofs):
            bfuns[itr]=np.zeros((ndofs,ndofs))
            for jtr in range(ndofs):
                bfuns[itr][ps[jtr][0],ps[jtr][1]]=V[itr,jtr]
        #print V.T
        for itr in range(ndofs):
            print bfuns[itr]

class ElementP1(ElementH1):

    n_dofs=1
    maxdeg=1

    def __init__(self,dim=1):
        self.dim=dim

    def lbasis(self,X,i):
        # TODO implement for 1D

        if self.dim==2:
            phi={
                0:lambda x,y: 1-x-y,
                1:lambda x,y: x,
                2:lambda x,y: y
                }[i](X[0],X[1])

            dphi={}
            dphi[0]={
                    0:lambda x,y: -1+0*x,
                    1:lambda x,y: 1+0*x,
                    2:lambda x,y: 0*x
                    }[i](X[0],X[1])
            dphi[1]={
                    0:lambda x,y: -1+0*x,
                    1:lambda x,y: 0*x,
                    2:lambda x,y: 1+0*x
                    }[i](X[0],X[1])
        elif self.dim==3:
            phi={
                0:lambda x,y,z: 1-x-y-z,
                1:lambda x,y,z: x,
                2:lambda x,y,z: y,
                3:lambda x,y,z: z,
                }[i](X[0],X[1],X[2])

            dphi={}
            dphi[0]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 1+0*x,
                    2:lambda x,y,z: 0*x,
                    3:lambda x,y,z: 0*x
                    }[i](X[0],X[1],X[2])
            dphi[1]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 0*x,
                    2:lambda x,y,z: 1+0*x,
                    3:lambda x,y,z: 0*x
                    }[i](X[0],X[1],X[2])
            dphi[2]={
                    0:lambda x,y,z: -1+0*x,
                    1:lambda x,y,z: 0*x,
                    2:lambda x,y,z: 0*x,
                    3:lambda x,y,z: 1+0*x
                    }[i](X[0],X[1],X[2])
        else:
            raise NotImplementedError("ElementP1.lbasis: not implemented for the given X.shape[0].")

        return phi,dphi

<|endoftext|>