<commit_before>"""
HAdmin is an API and command line tool to configure the queues of Hadoop's
CapacityScheduler.
"""

from argparse import ArgumentParser
import xml.etree.ElementTree as ET
import pkgutil
import os
import subprocess


scheduler_fname = 'capacity-scheduler.xml'
pre_scheduler = 'yarn.scheduler.capacity'
post_users = 'acl_submit_applications'
post_admins = 'acl_administer_queue'
post_cap = 'capacity'
post_maxcap = 'maximum-capacity'
post_state = 'state'
post_ulim = 'user-limit-factor'
post_subs = 'queues'


def queue_fqn(queue='root'):
    """ Returns the fully qualified name of a queue. """
    if (queue[0:4] == 'root'):
        return queue
    return ''.join(['root.', queue])

def queue_parent(queue='root'):
    """ Returns the parent of a queue. """
    return '.'.join(queue_fqn(queue).split('.')[0:-1])

def queue_cap_fqn(queue='root'):
    """ Returns the config key specifying the capacity of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_cap])

def queue_maxcap_fqn(queue='root'):
    """ Returns the config key specifying the max capacity of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_maxcap])

def queue_users_fqn(queue='root'):
    """ Returns the config key specifying the users of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_users])

def queue_admins_fqn(queue='root'):
    """ Returns the config key specifying the admins of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_admins])

def queue_state_fqn(queue='root'):
    """ Returns the config key specifying the state of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_state])

def queue_ulim_fqn(queue='root'):
    """ Returns the config key specifying the ulim of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_ulim])

def queue_subs_fqn(queue='root'):
    """ Returns the config key specifying the subqueues of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_subs])


class HXML:
    """
    A wrapper around etree specific to Hadoop's XML format.

    Provides an easy way to deal with Hadoop's XML so that calling classes
    only worry about Hadoop config keys and values, and not dealing with all
    the cruft. Unless you're crazy, avoid using this class directly and instead
    use a class such as QueueManager to fulfill your needs.
    """

    def __init__(self, etree):
        self.tree = etree

    @classmethod
    def from_etree(cls, etree):
        """ Not sure why this is here since __init__() takes etrees. """
        ret = cls(etree)
        return ret

    @classmethod
    def from_str(cls, string):
        """ Construct from a raw XML string. """
        ret = cls(ET.fromstring(string))
        return ret

    @classmethod
    def from_file(cls, fname):
        """ Construct from an XML file. """
        ret = cls(ET.parse(fname).getroot())
        return ret

    def __getitem__(self, prop):
        """ Retrieves a config value. """
        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                return node.find('value').text

        raise KeyError('Key ' + prop + ' not found')

    def __setitem__(self, prop, val):
        """
        Sets a config value. Please note that if the specified property
        prop does not exist, it will be created.
        """

        success = False

        if type(val) in (int, float):
            val = str(val)
        elif type(val) is bool:
            val = str(val).lower()

        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                node.find('value').text = val
                success = True

        if success is not True:
            el = ET.Element('property')
            el.text = "\n    "
            el.tail = "\n"
            name_el = ET.Element('name')
            name_el.tail = "\n    "
            name_el.text = prop
            el.append(name_el)
            val_el = ET.Element('value')
            val_el.tail = "\n  "
            val_el.text = val
            el.append(val_el)
            self.tree.append(el)

    def remove(self, prop):
        """ Deletes a property from the etree. """
        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                self.tree.remove(node)

    def keys(self):
        """ Returns a list of keys. """
        ret = list()
        for node in self.tree.findall('property'):
            ret.append(node.find('name').text)
        return sorted(ret)

    def save(self, fname):
        """ Saves to a file. """
        ET.ElementTree(self.tree).write(fname)


class QueueManager:
    """
    A class to specifically manage the queue's in Hadoop's CapacityScheduler.

    Contains a reasonable amount of functionality, and hopefully most of it
    is pretty common-sensical.
    """

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

    def __insert(self, prop, value):
        """ Inserts a value into a CSV string, checking for duplicates. """
        csv = [value]
        try:
            csv = self.hxml[prop].split(',')
            if value not in csv:
                csv.append(value)
        except KeyError:
            pass

        self.hxml[prop] = ','.join(sorted(csv))

    def __unsert(self, prop, value):
        """
        Unserts a value from a CSV string. I named this before realizing
        that 'remove' is actually the correct word. Good thing English is my
        only language.
        """

        csv = self.hxml[prop].split(',')

        if len(csv) == 1:
            raise ValueError("Can't delete last element in " + prop)

        for i in range(len(csv)):
            if csv[i] == value:
                del(csv[i])
                break

        self.hxml[prop] = ','.join(sorted(csv))

    def add_user(self, user, queue):
        """ Adds a user to a queue. """
        for user in user.split(','):
            self.__insert(queue_users_fqn(queue), user)

    def add_admin(self, user, queue):
        """ Adds an admin to a queue. """
        for user in user.split(','):
            self.__insert(queue_admins_fqn(queue), user)

    def set_cap(self, queue, cap):
        """ Sets a queue's capacity. """
        cap = int(cap)
        if cap > 100 or cap < 0:
            raise ValueError('Queue capacity must be between 0 and 100')

        self.hxml[queue_cap_fqn(queue)] = cap

    def set_maxcap(self, queue, maxcap):
        """ Set's a queue's maximum capacity. """
        maxcap = int(maxcap)
        if maxcap > 100 or maxcap < 0:
            raise ValueError(
                    'Queue maximum capacity must be between 0 and 100'
                    )

        self.hxml[queue_maxcap_fqn(queue)] = maxcap

    def set_ulim(self, queue, ulim):
        """ Set's a queue's user limit factor. """
        ulim = float(ulim)

        if ulim < 0.0:
            raise ValueError('ulim must be a positive float')

        self.hxml[queue_ulim_fqn(queue)] = ulim

    def set_state(self, queue, state):
        """ Set's a queue's state, which is either 'running' or 'stopped'. """
        self.hxml[queue_state_fqn(queue)] = state

    def off(self, queue):
        """ Convenience call to set a queue's state to 'stopped'. """
        self.set_state(queue, 'stopped')
<commit_msg>Parameters to 'queuecap' can now be in any order.<commit_after>"""
HAdmin is an API and command line tool to configure the queues of Hadoop's
CapacityScheduler.
"""

from argparse import ArgumentParser
import xml.etree.ElementTree as ET
import pkgutil
import os
import subprocess


scheduler_fname = 'capacity-scheduler.xml'
pre_scheduler = 'yarn.scheduler.capacity'
post_users = 'acl_submit_applications'
post_admins = 'acl_administer_queue'
post_cap = 'capacity'
post_maxcap = 'maximum-capacity'
post_state = 'state'
post_ulim = 'user-limit-factor'
post_subs = 'queues'


def queue_fqn(queue='root'):
    """ Returns the fully qualified name of a queue. """
    if (queue[0:4] == 'root'):
        return queue
    return ''.join(['root.', queue])

def queue_parent(queue='root'):
    """ Returns the parent of a queue. """
    return '.'.join(queue_fqn(queue).split('.')[0:-1])

def queue_cap_fqn(queue='root'):
    """ Returns the config key specifying the capacity of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_cap])

def queue_maxcap_fqn(queue='root'):
    """ Returns the config key specifying the max capacity of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_maxcap])

def queue_users_fqn(queue='root'):
    """ Returns the config key specifying the users of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_users])

def queue_admins_fqn(queue='root'):
    """ Returns the config key specifying the admins of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_admins])

def queue_state_fqn(queue='root'):
    """ Returns the config key specifying the state of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_state])

def queue_ulim_fqn(queue='root'):
    """ Returns the config key specifying the ulim of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_ulim])

def queue_subs_fqn(queue='root'):
    """ Returns the config key specifying the subqueues of a queue. """
    return '.'.join([pre_scheduler, queue_fqn(queue), post_subs])


class HXML:
    """
    A wrapper around etree specific to Hadoop's XML format.

    Provides an easy way to deal with Hadoop's XML so that calling classes
    only worry about Hadoop config keys and values, and not dealing with all
    the cruft. Unless you're crazy, avoid using this class directly and instead
    use a class such as QueueManager to fulfill your needs.
    """

    def __init__(self, etree):
        self.tree = etree

    @classmethod
    def from_etree(cls, etree):
        """ Not sure why this is here since __init__() takes etrees. """
        ret = cls(etree)
        return ret

    @classmethod
    def from_str(cls, string):
        """ Construct from a raw XML string. """
        ret = cls(ET.fromstring(string))
        return ret

    @classmethod
    def from_file(cls, fname):
        """ Construct from an XML file. """
        ret = cls(ET.parse(fname).getroot())
        return ret

    def __getitem__(self, prop):
        """ Retrieves a config value. """
        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                return node.find('value').text

        raise KeyError('Key ' + prop + ' not found')

    def __setitem__(self, prop, val):
        """
        Sets a config value. Please note that if the specified property
        prop does not exist, it will be created.
        """

        success = False

        if type(val) in (int, float):
            val = str(val)
        elif type(val) is bool:
            val = str(val).lower()

        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                node.find('value').text = val
                success = True

        if success is not True:
            el = ET.Element('property')
            el.text = "\n    "
            el.tail = "\n"
            name_el = ET.Element('name')
            name_el.tail = "\n    "
            name_el.text = prop
            el.append(name_el)
            val_el = ET.Element('value')
            val_el.tail = "\n  "
            val_el.text = val
            el.append(val_el)
            self.tree.append(el)

    def remove(self, prop):
        """ Deletes a property from the etree. """
        for node in self.tree.findall('property'):
            if node.find('name').text == prop:
                self.tree.remove(node)

    def keys(self):
        """ Returns a list of keys. """
        ret = list()
        for node in self.tree.findall('property'):
            ret.append(node.find('name').text)
        return sorted(ret)

    def save(self, fname):
        """ Saves to a file. """
        ET.ElementTree(self.tree).write(fname)


class QueueManager:
    """
    A class to specifically manage the queue's in Hadoop's CapacityScheduler.

    Contains a reasonable amount of functionality, and hopefully most of it
    is pretty common-sensical.
    """

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

    def __insert(self, prop, value):
        """ Inserts a value into a CSV string, checking for duplicates. """
        csv = [value]
        try:
            csv = self.hxml[prop].split(',')
            if value not in csv:
                csv.append(value)
        except KeyError:
            pass

        self.hxml[prop] = ','.join(sorted(csv))

    def __unsert(self, prop, value):
        """
        Unserts a value from a CSV string. I named this before realizing
        that 'remove' is actually the correct word. Good thing English is my
        only language.
        """

        csv = self.hxml[prop].split(',')

        if len(csv) == 1:
            raise ValueError("Can't delete last element in " + prop)

        for i in range(len(csv)):
            if csv[i] == value:
                del(csv[i])
                break

        self.hxml[prop] = ','.join(sorted(csv))

    def add_user(self, user, queue):
        """ Adds a user to a queue. """
        for user in user.split(','):
            self.__insert(queue_users_fqn(queue), user)

    def add_admin(self, user, queue):
        """ Adds an admin to a queue. """
        for user in user.split(','):
            self.__insert(queue_admins_fqn(queue), user)

    def set_cap(self, queue, cap):
        """ Sets a queue's capacity. """
        try:
            cap = int(cap)
        except ValueError:
            tmp = cap
            cap = int(queue)
            queue = tmp

        if cap > 100 or cap < 0:
            raise ValueError('Queue capacity must be between 0 and 100')

        self.hxml[queue_cap_fqn(queue)] = cap

    def set_maxcap(self, queue, maxcap):
        """ Set's a queue's maximum capacity. """
        try:
            maxcap = int(maxcap)
        except ValueError:
            tmp = maxcap
            maxcap = int(queue)
            queue = tmp

        if maxcap > 100 or maxcap < 0:
            raise ValueError(
                    'Queue maximum capacity must be between 0 and 100'
                    )

        self.hxml[queue_maxcap_fqn(queue)] = maxcap

    def set_ulim(self, queue, ulim):
        """ Set's a queue's user limit factor. """
        ulim = float(ulim)

        if ulim < 0.0:
            raise ValueError('ulim must be a positive float')

        self.hxml[queue_ulim_fqn(queue)] = ulim

    def set_state(self, queue, state):
        """ Set's a queue's state, which is either 'running' or 'stopped'. """
        self.hxml[queue_state_fqn(queue)] = state

    def off(self, queue):
        """ Convenience call to set a queue's state to 'stopped'. """
        self.set_state(queue, 'stopped')

{'fraction_non_alphanumeric': 0.12112171837708831, 'fraction_numerical': 0.003049588968443384, 'mean_word_length': 2.27081526452732, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 26}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1582177', 'n_tokens_mistral': 4694, 'n_tokens_neox': 4441, 'n_words': 1410}