<commit_before>from __future__ import print_function
from __future__ import unicode_literals
import time
import re
from netmiko.cisco_base_connection import CiscoBaseConnection
from netmiko.ssh_exception import NetMikoAuthenticationException
from netmiko import log

class HuaweiBase(CiscoBaseConnection):
    def session_preparation(self):
        """Prepare the session after the connection has been established."""
        self._test_channel_read()
        self.set_base_prompt()
        self.disable_paging(command="screen-length 0 temporary")
        # Clear the read buffer
        time.sleep(0.3 * self.global_delay_factor)
        self.clear_buffer()

    def config_mode(self, config_command="system-view"):
        """Enter configuration mode."""
        return super(HuaweiBase, self).config_mode(config_command=config_command)

    def exit_config_mode(self, exit_config="return", pattern=r">"):
        """Exit configuration mode."""
        return super(HuaweiBase, self).exit_config_mode(
            exit_config=exit_config, pattern=pattern
        )

    def check_config_mode(self, check_string="]"):
        """Checks whether in configuration mode. Returns a boolean."""
        return super(HuaweiBase, self).check_config_mode(check_string=check_string)

    def check_enable_mode(self, *args, **kwargs):
        """Huawei has no enable mode."""
        pass

    def enable(self, *args, **kwargs):
        """Huawei has no enable mode."""
        return ""

    def exit_enable_mode(self, *args, **kwargs):
        """Huawei has no enable mode."""
        return ""

    def set_base_prompt(
        self, pri_prompt_terminator=">", alt_prompt_terminator="]", delay_factor=1
    ):
        """
        Sets self.base_prompt

        Used as delimiter for stripping of trailing prompt in output.

        Should be set to something that is general and applies in multiple contexts. For Comware
        this will be the router prompt with < > or [ ] stripped off.

        This will be set on logging in, but not when entering system-view
        """
        log.debug("In set_base_prompt")
        delay_factor = self.select_delay_factor(delay_factor)
        self.clear_buffer()
        self.write_channel(self.RETURN)
        time.sleep(0.5 * delay_factor)

        prompt = self.read_channel()
        prompt = self.normalize_linefeeds(prompt)

        # If multiple lines in the output take the last line
        prompt = prompt.split(self.RESPONSE_RETURN)[-1]
        prompt = prompt.strip()

        # Check that ends with a valid terminator character
        if not prompt[-1] in (pri_prompt_terminator, alt_prompt_terminator):
            raise ValueError("Router prompt not found: {0}".format(prompt))

        # Strip off any leading HRP_. characters for USGv5 HA
        prompt = re.sub(r"^HRP_.", "", prompt, flags=re.M)

        # Strip off leading and trailing terminator
        prompt = prompt[1:-1]
        prompt = prompt.strip()
        self.base_prompt = prompt
        log.debug("prompt: {0}".format(self.base_prompt))

        return self.base_prompt

    def save_config(self, cmd="save", confirm=False, confirm_response=""):
        """ Save Config for HuaweiSSH"""
        return super(HuaweiBase, self).save_config(
            cmd=cmd, confirm=confirm, confirm_response=confirm_response
        )

class HuaweiSSH(HuaweiBase):
    """Huawei SSH driver."""

    pass

class HuaweiTelnet(HuaweiBase):
    """Huawei Telnet driver."""

    def telnet_login(
        self,
        pri_prompt_terminator=r"]\s*$",
        alt_prompt_terminator=r">\s*$",
        username_pattern=r"(?:user:|username|login|user name)",
        pwd_pattern=r"assword",
        delay_factor=1,
        max_loops=20,
    ):
        """Telnet login for Huawei Devices"""

        delay_factor = self.select_delay_factor(delay_factor)
        password_change_prompt = re.escape("Change now? [Y/N]")
        combined_pattern = r"({}|{}|{})".format(
            pri_prompt_terminator, alt_prompt_terminator, password_change_prompt
        )

        output = ""
        return_msg = ""
        i = 1
        while i <= max_loops:
            try:
                output = self.read_channel()
                return_msg += output

                # Search for username pattern / send username
                if re.search(username_pattern, output, flags=re.I):
                    self.write_channel(self.username + self.TELNET_RETURN)
                    time.sleep(1 * delay_factor)
                    output = self.read_channel()
                    return_msg += output

                # Search for password pattern / send password
                if re.search(pwd_pattern, output, flags=re.I):
                    self.write_channel(self.password + self.TELNET_RETURN)
                    time.sleep(0.5 * delay_factor)
                    output = self.read_channel()
                    return_msg += output
                    if re.search(
                        pri_prompt_terminator, output, flags=re.M
                    ) or re.search(alt_prompt_terminator, output, flags=re.M):
                        return return_msg

                # Search for password change prompt, send "N"
                if re.search(password_change_prompt, output):
                    self.write_channel("N" + self.TELNET_RETURN)
                    output = self.read_until_pattern(pattern=combined_pattern)
                    return_msg += output

                # Check if proper data received
                if re.search(pri_prompt_terminator, output, flags=re.M) or re.search(
                    alt_prompt_terminator, output, flags=re.M
                ):
                    return return_msg

                self.write_channel(self.TELNET_RETURN)
                time.sleep(0.5 * delay_factor)
                i += 1

            except EOFError:
                self.remote_conn.close()
                msg = "Login failed: {}".format(self.host)
                raise NetMikoAuthenticationException(msg)

        # Last try to see if we already logged in
        self.write_channel(self.TELNET_RETURN)
        time.sleep(0.5 * delay_factor)
        output = self.read_channel()
        return_msg += output
        if re.search(pri_prompt_terminator, output, flags=re.M) or re.search(
            alt_prompt_terminator, output, flags=re.M
        ):
            return return_msg

        self.remote_conn.close()
        msg = "Login failed: {}".format(self.host)
        raise NetMikoAuthenticationException(msg)

class HuaweiVrpv8SSH(HuaweiSSH):
    def commit(self, comment="", delay_factor=1):
        """
        Commit the candidate configuration.

        Commit the entered configuration. Raise an error and return the failure
        if the commit fails.

        default:
           command_string = commit
        comment:
           command_string = commit comment <comment>

        """
        delay_factor = self.select_delay_factor(delay_factor)
        error_marker = "Failed to generate committed config"
        command_string = "commit"

        if comment:
            command_string += ' comment "{}"'.format(comment)

        output = self.config_mode()
        output += self.send_command_expect(
            command_string,
            strip_prompt=False,
            strip_command=False,
            delay_factor=delay_factor,
            expect_string=r"]",
        )
        output += self.exit_config_mode()

        if error_marker in output:
            raise ValueError(
                "Commit failed with following errors:\n\n{}".format(output)
            )
        return output

    def save_config(self, *args, **kwargs):
        """Not Implemented"""
        raise NotImplementedError
<commit_msg>Improve the handling of logging in<commit_after>from __future__ import print_function
from __future__ import unicode_literals
import time
import re
from netmiko.cisco_base_connection import CiscoBaseConnection
from netmiko.ssh_exception import NetMikoAuthenticationException
from netmiko import log

class HuaweiBase(CiscoBaseConnection):
    def session_preparation(self):
        """Prepare the session after the connection has been established."""
        self._test_channel_read()
        self.set_base_prompt()
        self.disable_paging(command="screen-length 0 temporary")
        # Clear the read buffer
        time.sleep(0.3 * self.global_delay_factor)
        self.clear_buffer()

    def config_mode(self, config_command="system-view"):
        """Enter configuration mode."""
        return super(HuaweiBase, self).config_mode(config_command=config_command)

    def exit_config_mode(self, exit_config="return", pattern=r">"):
        """Exit configuration mode."""
        return super(HuaweiBase, self).exit_config_mode(
            exit_config=exit_config, pattern=pattern
        )

    def check_config_mode(self, check_string="]"):
        """Checks whether in configuration mode. Returns a boolean."""
        return super(HuaweiBase, self).check_config_mode(check_string=check_string)

    def check_enable_mode(self, *args, **kwargs):
        """Huawei has no enable mode."""
        pass

    def enable(self, *args, **kwargs):
        """Huawei has no enable mode."""
        return ""

    def exit_enable_mode(self, *args, **kwargs):
        """Huawei has no enable mode."""
        return ""

    def set_base_prompt(
        self, pri_prompt_terminator=">", alt_prompt_terminator="]", delay_factor=1
    ):
        """
        Sets self.base_prompt

        Used as delimiter for stripping of trailing prompt in output.

        Should be set to something that is general and applies in multiple contexts. For Comware
        this will be the router prompt with < > or [ ] stripped off.

        This will be set on logging in, but not when entering system-view
        """
        log.debug("In set_base_prompt")
        delay_factor = self.select_delay_factor(delay_factor)
        self.clear_buffer()
        self.write_channel(self.RETURN)
        time.sleep(0.5 * delay_factor)

        prompt = self.read_channel()
        prompt = self.normalize_linefeeds(prompt)

        # If multiple lines in the output take the last line
        prompt = prompt.split(self.RESPONSE_RETURN)[-1]
        prompt = prompt.strip()

        # Check that ends with a valid terminator character
        if not prompt[-1] in (pri_prompt_terminator, alt_prompt_terminator):
            raise ValueError("Router prompt not found: {0}".format(prompt))

        # Strip off any leading HRP_. characters for USGv5 HA
        prompt = re.sub(r"^HRP_.", "", prompt, flags=re.M)

        # Strip off leading and trailing terminator
        prompt = prompt[1:-1]
        prompt = prompt.strip()
        self.base_prompt = prompt
        log.debug("prompt: {0}".format(self.base_prompt))

        return self.base_prompt

    def save_config(self, cmd="save", confirm=False, confirm_response=""):
        """ Save Config for HuaweiSSH"""
        return super(HuaweiBase, self).save_config(
            cmd=cmd, confirm=confirm, confirm_response=confirm_response
        )

class HuaweiSSH(HuaweiBase):
    """Huawei SSH driver."""

    pass

class HuaweiTelnet(HuaweiBase):
    """Huawei Telnet driver."""

    def telnet_login(
        self,
        pri_prompt_terminator=r"]\s*$",
        alt_prompt_terminator=r">\s*$",
        username_pattern=r"(?:user:|username|login|user name)",
        pwd_pattern=r"assword",
        delay_factor=1,
        max_loops=20,
    ):
        """Telnet login for Huawei Devices"""

        delay_factor = self.select_delay_factor(delay_factor)
        password_change_prompt = re.escape("Change now? [Y/N]")
        combined_pattern = r"({}|{}|{})".format(
            pri_prompt_terminator, alt_prompt_terminator, password_change_prompt
        )

        output = ""
        return_msg = ""
        i = 1
        while i <= max_loops:
            try:
                # Search for username pattern / send username
                output = self.read_until_pattern(pattern=username_pattern, re_flags=re.I)
                return_msg += output
                self.write_channel(self.username + self.TELNET_RETURN)

                # Search for password pattern / send password
                output = self.read_until_pattern(pattern=pwd_pattern, re_flags=re.I)
                return_msg += output
                self.write_channel(self.password + self.TELNET_RETURN)

                # Waiting for combined output
                output = self.read_until_pattern(pattern=combined_pattern)
                return_msg += output

                # If received prompt, login is completed
                if re.search(
                    pri_prompt_terminator, output, flags=re.M
                ) or re.search(alt_prompt_terminator, output, flags=re.M):
                    return return_msg

                # Search for password change prompt, send "N"
                log.debug(f"Searching for [{password_change_prompt}]... delay_factor: {delay_factor}")
                if re.search(password_change_prompt, output):
                    log.debug(f"Found [{password_change_prompt}]...")
                    self.write_channel("N" + self.TELNET_RETURN)
                    output = self.read_until_pattern(pattern=combined_pattern)
                    return_msg += output

                # Check if proper data received
                if re.search(pri_prompt_terminator, output, flags=re.M) or re.search(
                    alt_prompt_terminator, output, flags=re.M
                ):
                    return return_msg

                self.write_channel(self.TELNET_RETURN)
                time.sleep(0.5 * delay_factor)
                i += 1

            except EOFError:
                self.remote_conn.close()
                msg = "Login failed: {}".format(self.host)
                raise NetMikoAuthenticationException(msg)

        # Last try to see if we already logged in
        self.write_channel(self.TELNET_RETURN)
        time.sleep(0.5 * delay_factor)
        output = self.read_channel()
        return_msg += output
        if re.search(pri_prompt_terminator, output, flags=re.M) or re.search(
            alt_prompt_terminator, output, flags=re.M
        ):
            return return_msg

        self.remote_conn.close()
        msg = "Login failed: {}".format(self.host)
        raise NetMikoAuthenticationException(msg)

class HuaweiVrpv8SSH(HuaweiSSH):
    def commit(self, comment="", delay_factor=1):
        """
        Commit the candidate configuration.

        Commit the entered configuration. Raise an error and return the failure
        if the commit fails.

        default:
           command_string = commit
        comment:
           command_string = commit comment <comment>

        """
        delay_factor = self.select_delay_factor(delay_factor)
        error_marker = "Failed to generate committed config"
        command_string = "commit"

        if comment:
            command_string += ' comment "{}"'.format(comment)

        output = self.config_mode()
        output += self.send_command_expect(
            command_string,
            strip_prompt=False,
            strip_command=False,
            delay_factor=delay_factor,
            expect_string=r"]",
        )
        output += self.exit_config_mode()

        if error_marker in output:
            raise ValueError(
                "Commit failed with following errors:\n\n{}".format(output)
            )
        return output

    def save_config(self, *args, **kwargs):
        """Not Implemented"""
        raise NotImplementedError
<|endoftext|>


<commit_before># pypcapfile.savefile.py
"""
Core functionality for reading and parsing libpcap savefiles. This contains
the core classes pcap_packet and pcap_savefile, as well as the core function
load_savefile.
"""

import binascii
import ctypes
import struct
import sys

import pcapfile.linklayer as linklayer

from pcapfile.structs import __pcap_header__, pcap_packet

VERBOSE = False

def __TRACE__(msg, args=None):
    if VERBOSE:
        if args:
            print msg % args
        else:
            print msg

class pcap_savefile(object):
    """
    Represents a libpcap savefile. The packets member is a list of pcap_packet
    instances. The 'valid' member will be None for an uninitialised instance,
    False if the initial validation fails, or True if the instance has been
    successfully set up and the file has been parsed.
    """
    def __init__(self, header, packets=None):
        if not packets:
            packets = []
        self.header = header
        self.packets = packets
        self.valid = None
        self.byteorder = sys.byteorder

        if not self.__validate__():
            self.valid = False
        else:
            self.valid = True

        assert self.valid, 'Invalid savefile.'

    def __validate__(self):
        assert __validate_header__(self.header), "Invalid header."
        if not __validate_header__(self.header):
            return False

        # TODO: extended validation
        valid_packet = lambda pkt: (pkt is not None or
                                    pkt.issubclass(ctypes.Structure))
        if not 0 == len(self.packets):
            valid_packet = [valid_packet(pkt) for pkt in self.packets]
            assert False not in valid_packet, 'Invalid packets in savefile.'
            if False in valid_packet:
                return False

        return True

    def __repr__(self):
        string = '%s-endian capture file version %d.%d\n'
        string += 'snapshot length: %d\n'
        string += 'linklayer type: %s\nnumber of packets: %d\n'
        string = string % (self.header.byteorder, self.header.major,
                           self.header.minor, self.header.snaplen,
                           linklayer.lookup(self.header.ll_type),
                           len(self.packets))
        return string

def _load_savefile_header(file_h):
    """
Load and validate the header of a pcap file.
    """
    raw_savefile_header = file_h.read(24)

    # in case the capture file is not the same endianness as ours, we have to
    # use the correct byte order for the file header
    if raw_savefile_header[:4] == '\xa1\xb2\xc3\xd4':
        byte_order = 'big'
        unpacked = struct.unpack('>IhhIIII', raw_savefile_header)
    elif raw_savefile_header[:4] == '\xd4\xc3\xb2\xa1':
        byte_order = 'little'
        unpacked = struct.unpack('<IhhIIII', raw_savefile_header)
    else:
        raise Exception('Invalid pcap file.')

    (magic, major, minor, tz_off, ts_acc, snaplen, ll_type) = unpacked
    header = __pcap_header__(magic, major, minor, tz_off, ts_acc, snaplen,
                             ll_type, ctypes.c_char_p(byte_order))
    if not __validate_header__(header):
        raise Exception('invalid savefile header!')
    else:
        return header

def load_savefile(input_file, layers=0, verbose=False):
    """
    Parse a savefile as a pcap_savefile instance. Returns the savefile
    on success and None on failure. Verbose mode prints additional information
    about the file's processing. layers defines how many layers to descend and
    decode the packet. input_file should be a Python file object.
    """
    global VERBOSE
    old_verbose = VERBOSE
    VERBOSE = verbose

    __TRACE__('[+] attempting to load %s', (input_file.name,))

    header = _load_savefile_header(input_file)
    if __validate_header__(header):
        __TRACE__('[+] found valid header')
        packets = _load_packets(input_file, header, layers)
        __TRACE__('[+] loaded %d packets', (len(packets),))
        sfile = pcap_savefile(header, packets)
        __TRACE__('[+] finished loading savefile.')
    else:
        __TRACE__('[!] invalid savefile')
        sfile = None

    VERBOSE = old_verbose
    return sfile

def __validate_header__(header):
    if not type(header) == __pcap_header__:
        return False

    if not header.magic == 0xa1b2c3d4:
        if not header.magic == 0xd4c3b2a1:
            return False

    assert header.byteorder in ['little', 'big'], 'Invalid byte order.'

    # as of savefile format 2.4, 'a 4-byte time zone offset; this
    # is always 0'; the same is true of the timestamp accuracy.
    if not header.tz_off == 0:
        return False

    if not header.ts_acc == 0:
        return False

    return True

def _load_packets(file_h, header, layers=0):
    """
    Read packets from the capture file. Expects the file handle to point to
    the location immediately after the header (24 bytes).
    """
    pkts = []

    hdrp = ctypes.pointer(header)
    while True:
        pkt = _read_a_packet(file_h, hdrp, layers)
        if pkt:
            pkts.append(pkt)
        else:
            break

    return pkts

def _read_a_packet(file_h, hdrp, layers=0):
    """
    Reads the next individual packet from the capture file. Expects
    the file handle to be somewhere after the header, on the next
    per-packet header.
    """
    raw_packet_header = file_h.read(16)
    if raw_packet_header == '':
        return None
    assert len(raw_packet_header) == 16, 'Unexpected end of per-packet header.'

    # in case the capture file is not the same endianness as ours, we have to
    # use the correct byte order for the packet header
    if hdrp[0].byteorder == 'big':
        packet_header = struct.unpack('>IIII', raw_packet_header)
    else:
        packet_header = struct.unpack('<IIII', raw_packet_header)
    (timestamp, timestamp_ms, capture_len, packet_len) = packet_header
    raw_packet_data = file_h.read(capture_len)

    # if the capture file is not the same endianness as ours, we need to
    # reverse the packet data
    if not __endian_check__(hdrp):
        raw_packet_data = raw_packet_data[::-1]
    assert len(raw_packet_data) == capture_len, 'Unexpected end of packet.'

    if layers > 0:
        layers -= 1
        raw_packet = linklayer.clookup(hdrp[0].ll_type)(raw_packet_data,
                                                        layers=layers)
    else:
        raw_packet = binascii.hexlify(raw_packet_data)

    packet = pcap_packet(hdrp, timestamp, timestamp_ms, capture_len,
                         packet_len, raw_packet)
    return packet

def __endian_check__(hdrp):
    if hdrp[0].magic == 0xa1b2c3d4:
        return True
    elif hdrp[0].magic == 0xd4c3b2a1:
        return False
    assert False, 'failed endian check.'
<commit_msg>Don't reverse the packet data.<commit_after># pypcapfile.savefile.py
"""
Core functionality for reading and parsing libpcap savefiles. This contains
the core classes pcap_packet and pcap_savefile, as well as the core function
load_savefile.
"""

import binascii
import ctypes
import struct
import sys

import pcapfile.linklayer as linklayer

from pcapfile.structs import __pcap_header__, pcap_packet

VERBOSE = False

def __TRACE__(msg, args=None):
    if VERBOSE:
        if args:
            print msg % args
        else:
            print msg

class pcap_savefile(object):
    """
    Represents a libpcap savefile. The packets member is a list of pcap_packet
    instances. The 'valid' member will be None for an uninitialised instance,
    False if the initial validation fails, or True if the instance has been
    successfully set up and the file has been parsed.
    """
    def __init__(self, header, packets=None):
        if not packets:
            packets = []
        self.header = header
        self.packets = packets
        self.valid = None
        self.byteorder = sys.byteorder

        if not self.__validate__():
            self.valid = False
        else:
            self.valid = True

        assert self.valid, 'Invalid savefile.'

    def __validate__(self):
        assert __validate_header__(self.header), "Invalid header."
        if not __validate_header__(self.header):
            return False

        # TODO: extended validation
        valid_packet = lambda pkt: (pkt is not None or
                                    pkt.issubclass(ctypes.Structure))
        if not 0 == len(self.packets):
            valid_packet = [valid_packet(pkt) for pkt in self.packets]
            assert False not in valid_packet, 'Invalid packets in savefile.'
            if False in valid_packet:
                return False

        return True

    def __repr__(self):
        string = '%s-endian capture file version %d.%d\n'
        string += 'snapshot length: %d\n'
        string += 'linklayer type: %s\nnumber of packets: %d\n'
        string = string % (self.header.byteorder, self.header.major,
                           self.header.minor, self.header.snaplen,
                           linklayer.lookup(self.header.ll_type),
                           len(self.packets))
        return string

def _load_savefile_header(file_h):
    """
Load and validate the header of a pcap file.
    """
    raw_savefile_header = file_h.read(24)

    # in case the capture file is not the same endianness as ours, we have to
    # use the correct byte order for the file header
    if raw_savefile_header[:4] == '\xa1\xb2\xc3\xd4':
        byte_order = 'big'
        unpacked = struct.unpack('>IhhIIII', raw_savefile_header)
    elif raw_savefile_header[:4] == '\xd4\xc3\xb2\xa1':
        byte_order = 'little'
        unpacked = struct.unpack('<IhhIIII', raw_savefile_header)
    else:
        raise Exception('Invalid pcap file.')

    (magic, major, minor, tz_off, ts_acc, snaplen, ll_type) = unpacked
    header = __pcap_header__(magic, major, minor, tz_off, ts_acc, snaplen,
                             ll_type, ctypes.c_char_p(byte_order))
    if not __validate_header__(header):
        raise Exception('invalid savefile header!')
    else:
        return header

def load_savefile(input_file, layers=0, verbose=False):
    """
    Parse a savefile as a pcap_savefile instance. Returns the savefile
    on success and None on failure. Verbose mode prints additional information
    about the file's processing. layers defines how many layers to descend and
    decode the packet. input_file should be a Python file object.
    """
    global VERBOSE
    old_verbose = VERBOSE
    VERBOSE = verbose

    __TRACE__('[+] attempting to load %s', (input_file.name,))

    header = _load_savefile_header(input_file)
    if __validate_header__(header):
        __TRACE__('[+] found valid header')
        packets = _load_packets(input_file, header, layers)
        __TRACE__('[+] loaded %d packets', (len(packets),))
        sfile = pcap_savefile(header, packets)
        __TRACE__('[+] finished loading savefile.')
    else:
        __TRACE__('[!] invalid savefile')
        sfile = None

    VERBOSE = old_verbose
    return sfile

def __validate_header__(header):
    if not type(header) == __pcap_header__:
        return False

    if not header.magic == 0xa1b2c3d4:
        if not header.magic == 0xd4c3b2a1:
            return False

    assert header.byteorder in ['little', 'big'], 'Invalid byte order.'

    # as of savefile format 2.4, 'a 4-byte time zone offset; this
    # is always 0'; the same is true of the timestamp accuracy.
    if not header.tz_off == 0:
        return False

    if not header.ts_acc == 0:
        return False

    return True

def _load_packets(file_h, header, layers=0):
    """
    Read packets from the capture file. Expects the file handle to point to
    the location immediately after the header (24 bytes).
    """
    pkts = []

    hdrp = ctypes.pointer(header)
    while True:
        pkt = _read_a_packet(file_h, hdrp, layers)
        if pkt:
            pkts.append(pkt)
        else:
            break

    return pkts

def _read_a_packet(file_h, hdrp, layers=0):
    """
    Reads the next individual packet from the capture file. Expects
    the file handle to be somewhere after the header, on the next
    per-packet header.
    """
    raw_packet_header = file_h.read(16)
    if raw_packet_header == '':
        return None
    assert len(raw_packet_header) == 16, 'Unexpected end of per-packet header.'

    # in case the capture file is not the same endianness as ours, we have to
    # use the correct byte order for the packet header
    if hdrp[0].byteorder == 'big':
        packet_header = struct.unpack('>IIII', raw_packet_header)
    else:
        packet_header = struct.unpack('<IIII', raw_packet_header)
    (timestamp, timestamp_ms, capture_len, packet_len) = packet_header
    raw_packet_data = file_h.read(capture_len)

    assert len(raw_packet_data) == capture_len, 'Unexpected end of packet.'

    if layers > 0:
        layers -= 1
        raw_packet = linklayer.clookup(hdrp[0].ll_type)(raw_packet_data,
                                                        layers=layers)
    else:
        raw_packet = binascii.hexlify(raw_packet_data)

    packet = pcap_packet(hdrp, timestamp, timestamp_ms, capture_len,
                         packet_len, raw_packet)
    return packet
<|endoftext|>


<commit_before>#!/usr/bin/python

# Copyright (c) 2009 Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Handle version information related to Visual Stuio."""

import os
import re
import subprocess
import sys

class VisualStudioVersion:
  """Information regarding a version of Visual Studio."""

  def __init__(self, short_name, description,
               solution_version, project_version, flat_sln):
    self.short_name = short_name
    self.description = description
    self.solution_version = solution_version
    self.project_version = project_version
    self.flat_sln = flat_sln

  def ShortName(self):
    return self.short_name

  def Description(self):
    """Get the full description of the version."""
    return self.description

  def SolutionVersion(self):
    """Get the version number of the sln files."""
    return self.solution_version

  def ProjectVersion(self):
    """Get the version number of the vcproj files."""
    return self.project_version

  def FlatSolution(self):
    return self.flat_sln

def _RegistryGetValue(key, value):
  """Use reg.exe to read a paricular key.

  While ideally we might use the win32 module, we would like gyp to be
  python neutral, so for instance cygwin python lacks this module.

  Arguments:
    key: The registry key to read from.
    value: The particular value to read.
  Return:
    The contents there, or None for failure.
  """
  # Skip if not on Windows.
  if sys.platform not in ('win32', 'cygwin'):
    return None
  # Run reg.exe.
  cmd = [os.path.join(os.environ.get('WINDIR', ''), 'System32', 'reg.exe'),
         'query', key, '/v', value]
  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  text = p.communicate()[0]
  # Require a successful return value.
  if p.returncode:
    return None
  # Extract value.
  match = re.search(r'REG_\w+[ ]+([^\r]+)\r\n', text)
  if not match:
    return None
  return match.group(1)

def _CreateVersion(name):
  versions = {
      '2008': VisualStudioVersion('2008',
                                  'Visual Studio 2008',
                                  solution_version='10.00',
                                  project_version='9.00',
                                  flat_sln=False),
      '2008e': VisualStudioVersion('2008e',
                                   'Visual Studio 2008',
                                   solution_version='10.00',
                                   project_version='9.00',
                                   flat_sln=True),
      '2005': VisualStudioVersion('2005',
                                  'Visual Studio 2005',
                                  solution_version='9.00',
                                  project_version='8.00',
                                  flat_sln=False),
      '2005e': VisualStudioVersion('2005e',
                                   'Visual Studio 2005',
                                   solution_version='9.00',
                                   project_version='8.00',
                                   flat_sln=True),
  }
  return versions[str(name)]

def _DetectVisualStudioVersions():
  """Collect the list of installed visual studio versions.

  Returns:
    A list of visual studio versions installed in descending order of
    usage preference.
    Base this on the registry and a quick check if devenv.exe exists.
    Only versions 8-9 are considered.
    Possibilities are:
      2005 - Visual Studio 2005 (8)
      2008 - Visual Studio 2008 (9)
  """
  version_to_year = {'8.0': '2005', '9.0': '2008'}
  versions = []
  for version in ('9.0', '8.0'):
    # Get the install dir for this version.
    key = r'HKLM\Software\Microsoft\VisualStudio\%s' % version
    path = _RegistryGetValue(key, 'InstallDir')
    if not path:
      continue
    # Check for full.
    if os.path.exists(os.path.join(path, 'devenv.exe')):
      # Add this one.
      versions.append(_CreateVersion(version_to_year[version]))
    # Check for express.
    elif os.path.exists(os.path.join(path, 'vcexpress.exe')):
      # Add this one.
      versions.append(_CreateVersion(version_to_year[version] + 'e'))
  return versions

def SelectVisualStudioVersion(version='auto'):
  """Select which version of Visual Studio projects to generate.

  Arguments:
    version: Hook to allow caller to force a particular version (vs auto).
  Returns:
    An object representing a visual studio project format version.
  """
  # In auto mode, check environment variable for override.
  if version == 'auto':
    version = os.environ.get('GYP_MSVS_VERSION', 'auto')
  # In auto mode, pick the most preferred version present.
  if version == 'auto':
    versions = _DetectVisualStudioVersions()
    if not versions:
      # Default to 2005.
      return _CreateVersion('2005')
    return versions[0]
  # Convert version string into a version object.
  return _CreateVersion(version)
<commit_msg>Fix detection of Visual Studio version.<commit_after>#!/usr/bin/python

# Copyright (c) 2009 Google Inc. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Handle version information related to Visual Stuio."""

import os
import re
import subprocess
import sys

class VisualStudioVersion:
  """Information regarding a version of Visual Studio."""

  def __init__(self, short_name, description,
               solution_version, project_version, flat_sln):
    self.short_name = short_name
    self.description = description
    self.solution_version = solution_version
    self.project_version = project_version
    self.flat_sln = flat_sln

  def ShortName(self):
    return self.short_name

  def Description(self):
    """Get the full description of the version."""
    return self.description

  def SolutionVersion(self):
    """Get the version number of the sln files."""
    return self.solution_version

  def ProjectVersion(self):
    """Get the version number of the vcproj files."""
    return self.project_version

  def FlatSolution(self):
    return self.flat_sln

def _RegistryGetValue(key, value):
  """Use reg.exe to read a paricular key.

  While ideally we might use the win32 module, we would like gyp to be
  python neutral, so for instance cygwin python lacks this module.

  Arguments:
    key: The registry key to read from.
    value: The particular value to read.
  Return:
    The contents there, or None for failure.
  """
  # Skip if not on Windows.
  if sys.platform not in ('win32', 'cygwin'):
    return None
  # Run reg.exe.
  cmd = [os.path.join(os.environ.get('WINDIR', ''), 'System32', 'reg.exe'),
         'query', key, '/v', value]
  p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  text = p.communicate()[0]
  # Require a successful return value.
  if p.returncode:
    return None
  # Extract value.
  match = re.search(r'REG_\w+\s+([^\r]+)\r\n', text)
  if not match:
    return None
  return match.group(1)

def _CreateVersion(name):
  versions = {
      '2008': VisualStudioVersion('2008',
                                  'Visual Studio 2008',
                                  solution_version='10.00',
                                  project_version='9.00',
                                  flat_sln=False),
      '2008e': VisualStudioVersion('2008e',
                                   'Visual Studio 2008',
                                   solution_version='10.00',
                                   project_version='9.00',
                                   flat_sln=True),
      '2005': VisualStudioVersion('2005',
                                  'Visual Studio 2005',
                                  solution_version='9.00',
                                  project_version='8.00',
                                  flat_sln=False),
      '2005e': VisualStudioVersion('2005e',
                                   'Visual Studio 2005',
                                   solution_version='9.00',
                                   project_version='8.00',
                                   flat_sln=True),
  }
  return versions[str(name)]

def _DetectVisualStudioVersions():
  """Collect the list of installed visual studio versions.

  Returns:
    A list of visual studio versions installed in descending order of
    usage preference.
    Base this on the registry and a quick check if devenv.exe exists.
    Only versions 8-9 are considered.
    Possibilities are:
      2005 - Visual Studio 2005 (8)
      2008 - Visual Studio 2008 (9)
  """
  version_to_year = {'8.0': '2005', '9.0': '2008'}
  versions = []
  for version in ('9.0', '8.0'):
    # Get the install dir for this version.
    key = r'HKLM\Software\Microsoft\VisualStudio\%s' % version
    path = _RegistryGetValue(key, 'InstallDir')
    if not path:
      continue
    # Check for full.
    if os.path.exists(os.path.join(path, 'devenv.exe')):
      # Add this one.
      versions.append(_CreateVersion(version_to_year[version]))
    # Check for express.
    elif os.path.exists(os.path.join(path, 'vcexpress.exe')):
      # Add this one.
      versions.append(_CreateVersion(version_to_year[version] + 'e'))
  return versions

def SelectVisualStudioVersion(version='auto'):
  """Select which version of Visual Studio projects to generate.

  Arguments:
    version: Hook to allow caller to force a particular version (vs auto).
  Returns:
    An object representing a visual studio project format version.
  """
  # In auto mode, check environment variable for override.
  if version == 'auto':
    version = os.environ.get('GYP_MSVS_VERSION', 'auto')
  # In auto mode, pick the most preferred version present.
  if version == 'auto':
    versions = _DetectVisualStudioVersions()
    if not versions:
      # Default to 2005.
      return _CreateVersion('2005')
    return versions[0]
  # Convert version string into a version object.
  return _CreateVersion(version)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Neuroimaging tools for Python (NIPY).

The aim of NIPY is to produce a platform-independent Python environment for
the analysis of brain imaging data using an open development model.

While
the project is still in its initial stages, packages for file I/O, script
support as well as single subject fMRI and random effects group comparisons
model are currently available.

Specifically, we aim to:

   1. Provide an open source, mixed language scientific programming
      environment suitable for rapid development.

   2. Create sofware components in this environment to make it easy
      to develop tools for MRI, EEG, PET and other modalities.

   3. Create and maintain a wide base of developers to contribute to
      this platform.

   4. To maintain and develop this framework as a single, easily
      installable bundle.

Package Organization
====================
The neuroimaging package contains the following subpackages and modules:

.. packagetree::
   :style: UML
"""
__docformat__ = 'restructuredtext en'

from version import version as __version__
# FIXME
#__revision__ = int("$Rev$".split()[-2])
__status__   = 'alpha'
__date__    = "$LastChangedDate$"
__url__     = 'http://neuroimaging.scipy.org'

packages = (
  'neuroimaging',
  'neuroimaging.algorithms',
  'neuroimaging.algorithms.tests',
  'neuroimaging.algorithms.statistics',
  'neuroimaging.algorithms.statistics.tests',
  'neuroimaging.core',
  'neuroimaging.core.image',
  'neuroimaging.core.image.tests',
  'neuroimaging.core.reference',
  'neuroimaging.core.reference.tests',
  'neuroimaging.io',
  'neuroimaging.io.tests',
  'neuroimaging.modalities',
  'neuroimaging.modalities.fmri',
  'neuroimaging.modalities.fmri.tests',
  'neuroimaging.modalities.fmri.fmristat',
  'neuroimaging.modalities.fmri.fmristat.tests',
  'neuroimaging.utils',
  'neuroimaging.utils.tests',
  'neuroimaging.utils.tests.data',
  'neuroimaging.testing')

def import_from(modulename, objectname):
    """Import and return objectname from modulename."""
    module = __import__(modulename, {}, {}, (objectname,))
    try:
        return getattr(module, objectname)
    except AttributeError:
        return None

from neuroimaging.testing import Tester

test = Tester().test
bench = Tester().bench

def _test_local_install():
    """ Warn the user that running with neuroimaging being
        imported locally is a bad idea.
    """
    import os
    if os.getcwd() == os.sep.join(
                            os.path.abspath(__file__).split(os.sep)[:-2]):
        import warnings
        warnings.warn('Running the tests from the install directory may '
                     'trigger some failures')

_test_local_install()
<commit_msg>BUG: Fail gracefully if user has old version of numpy without nose test framework.<commit_after># -*- coding: utf-8 -*-
"""
Neuroimaging tools for Python (NIPY).

The aim of NIPY is to produce a platform-independent Python environment for
the analysis of brain imaging data using an open development model.

While
the project is still in its initial stages, packages for file I/O, script
support as well as single subject fMRI and random effects group comparisons
model are currently available.

Specifically, we aim to:

   1. Provide an open source, mixed language scientific programming
      environment suitable for rapid development.

   2. Create sofware components in this environment to make it easy
      to develop tools for MRI, EEG, PET and other modalities.

   3. Create and maintain a wide base of developers to contribute to
      this platform.

   4. To maintain and develop this framework as a single, easily
      installable bundle.

Package Organization
====================
The neuroimaging package contains the following subpackages and modules:

.. packagetree::
   :style: UML
"""
__docformat__ = 'restructuredtext en'

from version import version as __version__
# FIXME
#__revision__ = int("$Rev$".split()[-2])
__status__   = 'alpha'
__date__    = "$LastChangedDate$"
__url__     = 'http://neuroimaging.scipy.org'

packages = (
  'neuroimaging',
  'neuroimaging.algorithms',
  'neuroimaging.algorithms.tests',
  'neuroimaging.algorithms.statistics',
  'neuroimaging.algorithms.statistics.tests',
  'neuroimaging.core',
  'neuroimaging.core.image',
  'neuroimaging.core.image.tests',
  'neuroimaging.core.reference',
  'neuroimaging.core.reference.tests',
  'neuroimaging.io',
  'neuroimaging.io.tests',
  'neuroimaging.modalities',
  'neuroimaging.modalities.fmri',
  'neuroimaging.modalities.fmri.tests',
  'neuroimaging.modalities.fmri.fmristat',
  'neuroimaging.modalities.fmri.fmristat.tests',
  'neuroimaging.utils',
  'neuroimaging.utils.tests',
  'neuroimaging.utils.tests.data',
  'neuroimaging.testing')

def import_from(modulename, objectname):
    """Import and return objectname from modulename."""
    module = __import__(modulename, {}, {}, (objectname,))
    try:
        return getattr(module, objectname)
    except AttributeError:
        return None

try:
    from neuroimaging.testing import Tester
    test = Tester().test
    bench = Tester().bench
except ImportError:
    # If the user has an older version of numpy which does not have
    # the nose test framework, fail gracefully and prompt them to
    # upgrade.
    import numpy as np
    npver = np.__version__.split('.')
    npver = '.'.join((npver[0], npver[1]))
    npver = float(npver)
    if npver < 1.2:
        raise ImportError('Nipy requires numpy version 1.2 or greater. '
                          '\n    You have numpy version %s installed.'
                          '\n    Please upgrade numpy:  '
                          'http://www.scipy.org/NumPy'
                          % np.__version__)

def _test_local_install():
    """ Warn the user that running with neuroimaging being
        imported locally is a bad idea.
    """
    import os
    if os.getcwd() == os.sep.join(
                            os.path.abspath(__file__).split(os.sep)[:-2]):
        import warnings
        warnings.warn('Running the tests from the install directory may '
                     'trigger some failures')

_test_local_install()
<|endoftext|>


<commit_before># !/usr/bin/env python

"""
This module implements representations of slabs and surfaces.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2014, The Materials Virtual Lab"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "6/10/14"

from fractions import gcd
import math
import numpy as np
import itertools

from pymatgen.core.structure import Structure

def lcm(numbers):
    """Return lowest common multiple."""
    def lcm(a, b):
        return (a * b) / gcd(a, b)
    return reduce(lcm, numbers, 1)

class Slab(Structure):
    """
    Subclass of Structure representing a Slab. Implements additional
    attributes pertaining to slabs.

    .. attribute:: parent

        Parent structure from which Slab was derived.

    .. attribute:: min_slab_size

        Minimum size in angstroms of layers containing atoms

    .. attribute:: min_vac_size

        Minimize size in angstroms of layers containing vacuum

    .. attribute:: scale_factor

        Final computed scale factor that brings the parent cell to the
        surface cell.

    .. attribute:: normal

        Surface normal vector.
    """

    def __init__(self, structure, miller_index, min_slab_size,
                 min_vacuum_size, lll_reduce=True, shift=0):
        """
        Makes a Slab structure. Note that the code will make a slab with
        whatever size that is specified, rounded upwards. The a and b lattice
        vectors are always in-plane, and the c lattice is always out of plane
        (though not necessarily orthogonal).

        Args:
            structure (Structure): Initial input structure.
            miller_index ([h, k, l]): Miller index of plane parallel to
                surface. Note that this is referenced to the input structure. If
                you need this to be based on the conventional cell,
                you should supply the conventional structure.
            min_slab_size (float): In Angstroms
            min_vacuum_size (float): In Angstroms
            lll_reduce (bool): Whether to perform an LLL reduction on the
                eventual structure.
            shift (float): In Angstroms (shifting the origin)
        """
        latt = structure.lattice
        d = reduce(gcd, miller_index)
        miller_index = [int(i / d) for i in miller_index]
        #Calculate the surface normal using the reciprocal lattice vector.
        recp = latt.reciprocal_lattice_crystallographic
        normal = recp.get_cartesian_coords(miller_index)
        normal /= np.linalg.norm(normal)

        slab_scale_factor = []
        non_orth_ind = []
        eye = np.eye(3, dtype=np.int)
        dist = float('inf')
        for i, j in enumerate(miller_index):
            if j == 0:
                # Lattice vector is perpendicular to surface normal, i.e.,
                # in plane of surface. We will simply choose this lattice
                # vector as one of the basis vectors.
                slab_scale_factor.append(eye[i])
            else:
                #Calculate projection of lattice vector onto surface normal.
                d = abs(np.dot(normal, latt.matrix[i]))
                non_orth_ind.append(i)
                if d < dist:
                    latt_index = i
                    dist = d

        if len(non_orth_ind) > 1:
            lcm_miller = lcm([miller_index[i] for i in non_orth_ind])
            for i, j in itertools.combinations(non_orth_ind, 2):
                l = [0, 0, 0]
                l[i] = -int(round(lcm_miller / miller_index[i]))
                l[j] = int(round(lcm_miller / miller_index[j]))
                slab_scale_factor.append(l)
                if len(slab_scale_factor) == 2:
                    break

        nlayers_slab = int(math.ceil(min_slab_size / dist))
        nlayers_vac = int(math.ceil(min_vacuum_size / dist))
        nlayers = nlayers_slab + nlayers_vac
        slab_scale_factor.append(eye[latt_index] * nlayers)

        slab = structure.copy()

        slab.make_supercell(slab_scale_factor)
        new_sites = []
        for site in slab:
            if shift <= np.dot(site.coords, normal) < nlayers_slab * dist + \
                    shift:
                new_sites.append(site)
        slab = Structure.from_sites(new_sites)

        if lll_reduce:
            lll_slab = slab.copy(sanitize=True)
            mapping = lll_slab.lattice.find_mapping(slab.lattice)
            slab_scale_factor = np.dot(mapping[2], slab_scale_factor)
            slab = lll_slab

        self.parent = structure
        self.min_slab_size = min_slab_size
        self.min_vac_size = min_vacuum_size
        self.scale_factor = np.array(slab_scale_factor)
        self.normal = normal
        self.miller_index = miller_index

        super(Slab, self).__init__(
            slab.lattice, slab.species_and_occu, slab.frac_coords,
            site_properties=slab.site_properties)

    @property
    def surface_area(self):
        m = self.lattice.matrix
        return np.linalg.norm(np.cross(m[0], m[1]))

    def add_adsorbate_atom(self, indices, specie, distance):
        """
        Gets the structure of single atom adsorption.
        slab structure from the Slab class(in [0, 0, 1])

        Args:
            indices ([int]): Indices of sites on which to put the absorbate.
                Absorbed atom will be displaced relative to the center of
                these sites.
            specie (Specie/Element/str): adsorbed atom species
            distance (float): between centers of the adsorbed atom and the
                given site in Angstroms.
        """
        #Let's do the work in cartesian coords
        center = np.sum([self[i].coords for i in indices], axis=0) / len(
            indices)

        coords = center + self.normal * distance / np.linalg.norm(self.normal)

        self.append(specie, coords, coords_are_cartesian=True)

import unittest
from pymatgen.core.lattice import Lattice

class SlabTest(unittest.TestCase):

    def setUp(self):
        self.cu = Structure(Lattice.cubic(3), ["Cu", "Cu", "Cu", "Cu"],
                            [[0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5],
                             [0, 0.5, 0.5]])

    def test_init(self):
        for hkl in itertools.product(xrange(4), xrange(4), xrange(4)):
            if any(hkl):
                ssize = 6
                vsize = 10
                s = Slab(self.cu, hkl, ssize, vsize)
                if hkl == [0, 1, 1]:
                    self.assertEqual(len(s), 13)
                    self.assertAlmostEqual(s.surface_area, 12.727922061357855)
                manual = self.cu.copy()
                manual.make_supercell(s.scale_factor)
                self.assertEqual(manual.lattice.lengths_and_angles,
                                 s.lattice.lengths_and_angles)

        # # For visual debugging
        # from pymatgen import write_structure
        # write_structure(s.parent, "cu.cif")
        # write_structure(s, "cu_slab_%s_%.3f_%.3f.cif" %
        #                  (str(hkl), ssize, vsize))

    def test_adsorp_atom(self):
        # O adsorb on 4Cu[0.5, 0.5, 0.25], abc = [3, 3, 12]
        # 1. test site_a = abc input
        s1 = Slab(self.cu, [0, 0, 1], 5, 5)
        s1.add_adsorbate_atom([3], 'O', 2)

        self.assertEqual(len(s1), 9)
        self.assertEqual(s1.formula, "Cu8 O1")
        self.assertTrue(s1.ntypesp == 2)
        self.assertTrue(s1.symbol_set == ("Cu", "O"))
        index1 = int(s1.indices_from_symbol('O')[0])
        self.assertAlmostEqual(s1[index1].a, 0.5)
        self.assertAlmostEqual(s1[index1].b, 0.5)
        self.assertAlmostEqual(s1[index1].c, 0.4166667)

if __name__ == "__main__":
    unittest.main()

<commit_msg>Update surface.py<commit_after># !/usr/bin/env python

"""
This module implements representations of slabs and surfaces.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2014, The Materials Virtual Lab"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "6/10/14"

from fractions import gcd
import math
import numpy as np
import itertools

from pymatgen.core.structure import Structure
from pymatgen import Lattice, Structure

import numpy as np
import copy
import scipy.cluster.hierarchy

def lcm(numbers):
    """Return lowest common multiple."""
    def lcm(a, b):
        return (a * b) / gcd(a, b)
    return reduce(lcm, numbers, 1)

class Slab(Structure):
    """
    Subclass of Structure representing a Slab. Implements additional
    attributes pertaining to slabs.

    .. attribute:: parent

        Parent structure from which Slab was derived.

    .. attribute:: min_slab_size

        Minimum size in angstroms of layers containing atoms

    .. attribute:: min_vac_size

        Minimize size in angstroms of layers containing vacuum

    .. attribute:: scale_factor

        Final computed scale factor that brings the parent cell to the
        surface cell.

    .. attribute:: normal

        Surface normal vector.
    """

    def __init__(self, structure, miller_index, min_slab_size,
                 min_vacuum_size, temp_thresh, lll_reduce=True, shift=0):
        """
        Makes a Slab structure. Note that the code will make a slab with
        whatever size that is specified, rounded upwards. The a and b lattice
        vectors are always in-plane, and the c lattice is always out of plane
        (though not necessarily orthogonal).

        Args:
            structure (Structure): Initial input structure.
            miller_index ([h, k, l]): Miller index of plane parallel to
                surface. Note that this is referenced to the input structure. If
                you need this to be based on the conventional cell,
                you should supply the conventional structure.
            min_slab_size (float): In Angstroms
            min_vacuum_size (float): In Angstroms
            temp_thresh (float): A temporary parameter used for setting the
                t parameter in fclusterdata in order to determine the number of
                terminations to be found. An algorithm wil eventually be implemented
                so that the program can determine this value by itself depending on
                the nature of the structure object parameter.
            lll_reduce (bool): Whether to perform an LLL reduction on the
                eventual structure.
            shift (float): In Angstroms (shifting the origin)
        """
        latt = structure.lattice
        d = reduce(gcd, miller_index)
        miller_index = [int(i / d) for i in miller_index]
        #Calculate the surface normal using the reciprocal lattice vector.
        recp = latt.reciprocal_lattice_crystallographic
        normal = recp.get_cartesian_coords(miller_index)
        normal /= np.linalg.norm(normal)

        slab_scale_factor = []
        non_orth_ind = []
        eye = np.eye(3, dtype=np.int)
        dist = float('inf')
        for i, j in enumerate(miller_index):
            if j == 0:
                # Lattice vector is perpendicular to surface normal, i.e.,
                # in plane of surface. We will simply choose this lattice
                # vector as one of the basis vectors.
                slab_scale_factor.append(eye[i])
            else:
                #Calculate projection of lattice vector onto surface normal.
                d = abs(np.dot(normal, latt.matrix[i]))
                non_orth_ind.append(i)
                if d < dist:
                    latt_index = i
                    dist = d

        if len(non_orth_ind) > 1:
            lcm_miller = lcm([miller_index[i] for i in non_orth_ind])
            for i, j in itertools.combinations(non_orth_ind, 2):
                l = [0, 0, 0]
                l[i] = -int(round(lcm_miller / miller_index[i]))
                l[j] = int(round(lcm_miller / miller_index[j]))
                slab_scale_factor.append(l)
                if len(slab_scale_factor) == 2:
                    break

        min_slab_size = float(min_slab_size)
        min_vacuum_size = float(min_vacuum_size)
        nlayers_slab = int(math.ceil(min_slab_size / dist))
        nlayers_vac = int(math.ceil(min_vacuum_size / dist))
        nlayers = nlayers_slab + nlayers_vac

        slab_scale_factor.append(eye[latt_index] * nlayers)
        slab = structure.copy()
        slab.make_supercell(slab_scale_factor)
        new_sites = []

        # This for loop needs to take into account fractional atoms in a unit cell when c = 0 or c = 1 .
        # The technique to do so is highly volatile, make sure to analyze it later if something goes wrong. Such as when
        # a termination is missing from one of the faces of the cell.
        for site in slab:
            if shift <= np.dot(site.coords, normal) <= min_slab_size + shift:
                new_sites.append(site)
            if shift == 0 and np.dot(site.coords, normal) <= 0:
                new_sites.append(site)
            if shift == 0 and (site.coords[2] == slab.lattice.c or site.coords[1] == slab.lattice.c or site.coords[0] == slab.lattice.c):
                new_sites.append(site)

        slab = Structure.from_sites(new_sites)

# After this line is where all the necessary algorithms need to enumerate the different terminations is inserted
        el = slab.species
        org_coords = slab.frac_coords.tolist()
        new_coord, b = [], []

        for i in slab.frac_coords:
            b.append(i[2])
            new_coord.append(b)
            b = []

        tracker_index = scipy.cluster.hierarchy.fclusterdata(new_coord, temp_thresh, criterion= 'distance')
        new_coord, tracker_index, org_coords, el = zip(*sorted(zip(new_coord, tracker_index, org_coords, el)))

# Creates a list that keeps track of repetitive values of c in different coordinates
        org_index = []
        for i in range(0, len(slab)):
            if i == len(slab) - 1:
                org_index.append(i)
            else:
                if tracker_index[i] != tracker_index[i+1]:
                    org_index.append(i)

        term_enum, term_coords = [], []
        a, i = 0, 0
        new_slab = Structure(slab.lattice, el, org_coords)

        for iii in org_index:
            y = []
            alt_slab = new_slab.copy()
            n = 0

            for ii in range(0, len(alt_slab)):
                index = []
                term_scale = 0
                index.append(ii)
                if alt_slab.frac_coords[ii][2] > alt_slab.frac_coords[iii][2]:
                    term_scale = min_vacuum_size/(dist*nlayers)
                if float(alt_slab.frac_coords[ii][2]) + term_scale > 1: # Out of bounds atom
                    print("Warning, the shift in the origin may cause some sites be pushed outside the slab")
                    continue
                alt_slab.translate_sites(index, [0, 0, term_scale])
                n += 1

            term_enum.append(alt_slab)

            for iv in range(a, iii + 1):
                    y.append(term_enum[i][iv])
            i += 1
            term_coords.append(y)
            a = iii+1

        if lll_reduce:
            lll_slab = slab.copy(sanitize=True)
            mapping = lll_slab.lattice.find_mapping(slab.lattice)
            slab_scale_factor = np.dot(mapping[2], slab_scale_factor)
            slab = lll_slab

        self.min_slab_size = min_slab_size
        self.min_vac_size = min_vacuum_size
        self.enum = term_enum # Holds a list of Structure objects of slabs with different terminations
        self.parent = structure
        self.miller_index = miller_index
        self.org_in = org_index
        self.term_coords = term_coords # Holds the corresponding list of sites on the surface terminations
        self.temp_thresh = temp_thresh

        super(Slab, self).__init__(
            slab.lattice, slab.species_and_occu, slab.frac_coords,
            site_properties=slab.site_properties)

    @property
    def surface_area(self):
        m = self.lattice.matrix
        return np.linalg.norm(np.cross(m[0], m[1]))

    @classmethod
    def adsorb_atom(cls, structure_a, site_a, atom, distance,
                    surface=[0, 0, 1], xyz=0):
        """
        Gets the structure of single atom adsorption.

        Args:
        structure_a: the slab structure for adsorption
        site_a:  given sites for adsorption.
             default(xyz=0): site_a = [a, b, c], within [0,1];
             xyz=1: site_a = [x, y, z], in Angstroms.
        atom: adsorbed atom species
        distance: between centers of the adsorbed atom and the given site.
             in Angstroms
        surface: direction of the surface where atoms are adsorbed to
             default: surface = [0, 0, 1]
        xyz: default is 0. 1 means site_a = [x, y, z], in Angstroms.

        """
        from pymatgen.transformations.site_transformations import \
            InsertSitesTransformation

        lattice_s = structure_a.lattice
        abc_s = lattice_s.abc
        # a123_s = lattice_s.matrix
        b123_s = lattice_s.inv_matrix.T
        # print surface
        vector_o = np.dot(surface, b123_s)
        print vector_o
        lens_v = np.sqrt(np.dot(vector_o, vector_o.T))
        V_o = vector_o / lens_v * distance

        if xyz == 0:
            # site_a = [a, b, c]
            for i in xrange(3):
                if site_a[i]> 1 or site_a[i] < 0:
                    raise ValueError("site_a is outsite the cell.")
            site_abc = V_o / abc_s + site_a
        else:
            # site_a = [x, y, z]
            for i in xrange(3):
                if site_a[i] > abc_s[i]:
                    raise ValueError("sites_a is outside the cell.")
            site_a1 = np.array(site_a)

            # convert to a1, a2, a3
            #< site_a2 = np.dot(a123_s, site_a1.T) / abc_s
            #< site_abc = (V_o+site_a2) / abc_s
            site_a2 = np.dot(b123_s, site_a1.T)

            site_abc = V_o/abc_s+site_a2

        for i in xrange(3):
            if site_abc[i] < 0 or site_abc[i] > 1:
                raise ValueError("wrong distance, atom will be outside the cell.")

        print 'add_site:', site_abc, atom

        ist = InsertSitesTransformation(species=atom, coords=[site_abc])
        structure_ad = ist.apply_transformation(structure_a)

        return structure_ad

from pymatgen.io.smartio import CifParser
from pymatgen import write_structure
import unittest

class SlabTest(unittest.TestCase): # 4 different unit cells will be tested, ZnO, BaTiO4, TeI, and LifePO4

    def setUp(self):

        Zn_O = CifParser("/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/ZnO-wz.cif")
        self.zno = (Zn_O.get_structures(primitive = False)[0])
        Te_I  = CifParser("/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/icsd_TeI.cif")
        self.tei= (Te_I.get_structures(primitive = False)[0])
        Ba_Ti_O3  = CifParser("/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/icsd_batio3.cif")
        self.batio3= (Ba_Ti_O3.get_structures(primitive = False)[0])
        Li_Fe_PO4  = CifParser("/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/LiFePO4.cif")
        self.lifepo4= (Li_Fe_PO4.get_structures(primitive = False)[0])

    def test_init(self):

        hkl001 = [0, 0, 1]
        z001 = Slab(self.zno, hkl001, 10, 3, 0.025, shift = 2)
        t001 = Slab(self.tei, hkl001, 10, 3, 0.025, shift = 5)
        b001 = Slab(self.batio3, hkl001, 10, 3, 0.01, shift = 2)
        l001 = Slab(self.lifepo4, hkl001, 10, 3, 0.01, shift = 2)

        hkl100 = [1, 0, 0]
        z100 = Slab(self.zno, hkl100, 10, 5, 0.01)
        t100 = Slab(self.tei, hkl100, 10, 5, 0.01)
        b100 = Slab(self.batio3, hkl100, 10, 5, 0.01, shift = 3)
        l100 = Slab(self.lifepo4, hkl100, 10, 3, 0.01, shift = 2)

        m = [z001, t001, b001, l001, z100, t100, b100, l100]

        fileName = ["/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/ZnO-wz.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/icsd_TeI.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/icsd_batio3.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/001_terminations/LiFePO4.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/100_terminations/ZnO-wz.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/100_terminations/icsd_TeI.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/100_terminations/icsd_batio3.cif",
                    "/home/richard/pymatgen/pymatgen/core/surface tests/100_terminations/LiFePO4.cif"]

        for i in range(0, len(fileName)):
            Name = fileName[i][:-4] + "_%s_%s_%s_%s_" \
                                      %(str(m[i].miller_index),
                                        m[i].min_slab_size,
                                        m[i].min_vac_size,
                                        m[i].temp_thresh)
            fileType = ".cif"

            # For visual debugging
            for iii in range(0, len(m[i].enum)):
                name_num = str(iii)
                newFile = Name + name_num + fileType
                write_structure(m[i].enum[iii], newFile)

                # Compares the newly created structure to one that was already made
                test_comp = CifParser("/home/richard/pymatgen/pymatgen/core/surface tests/tests/" + newFile[+68:])
                # Line 347 might change depending on where you cloned pymatgen
                self.compare_structure = (test_comp.get_structures(primitive = False)[0])
                test_new = CifParser(newFile)
                self.new_structure = (test_new.get_structures(primitive = False)[0])
                self.assertEqual(self.compare_structure, self.new_structure) #This may take a while to process

            # Prints the coordinates and the species of each atom along with the number of atoms on a
            # surface termination site
            for iiii in range(0, len(m[i].term_coords)):
                print(m[i].term_coords[iiii])
                print("%s atoms in this termination surface." %(len(m[i].term_coords[iiii])))

            print(" ")
        # Checks to see if the program generates the number of terminations we would expect it to
        self.assertEqual(len(z001.enum), 7)
        self.assertEqual(len(b001.enum), 5)
        self.assertEqual(len(t001.enum), 8)
        self.assertEqual(len(z100.enum), 7)
        self.assertEqual(len(b100.enum), 5)
        self.assertEqual(len(t100.enum), 8)

if __name__ == "__main__":
    unittest.main()

class SlabTest(unittest.TestCase):

    def setUp(self):
        self.cu = Structure(Lattice.cubic(3), ["Cu", "Cu", "Cu", "Cu"],
                            [[0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5],
                             [0, 0.5, 0.5]])

    def test_init(self):
        for hkl in itertools.product(xrange(4), xrange(4), xrange(4)):
            if any(hkl):
                ssize = 6
                vsize = 10
                s = Slab(self.cu, hkl, ssize, vsize)
                if hkl == [0, 1, 1]:
                    self.assertEqual(len(s), 13)
                    self.assertAlmostEqual(s.surface_area, 12.727922061357855)
                manual = self.cu.copy()
                manual.make_supercell(s.scale_factor)
                self.assertEqual(manual.lattice.lengths_and_angles,
                                 s.lattice.lengths_and_angles)

         # For visual debugging
        from pymatgen import write_structure
        write_structure(s.parent, "cu.cif")
        write_structure(s, "cu_slab_%s_%.3f_%.3f.cif" %
                          (str(hkl), ssize, vsize))

    def test_adsorb_atom(self):
        s001 = Slab(self.cu,[0, 0, 1], 5, 5)
        print s001
        #O adsorb on 4Cu[0.5, 0.5, 0.25], abc = [3, 3, 12]
         #1. test site_a = abc input
        s001_ad1 = Slab.adsorb_atom(structure_a=s001, site_a=[0.5, 0.5, 0.25], atom= ['O'],
                                    distance=2)
        self.assertEqual(len(s001_ad1), 9)
        for i in xrange(len(s001_ad1)):
            if str(s001_ad1[i].specie) == 'O':
                print s001_ad1[i].frac_coords
                self.assertAlmostEqual(s001_ad1[i].a, 0.5)
                self.assertAlmostEqual(s001_ad1[i].b, 0.5)
                self.assertAlmostEqual(s001_ad1[i].c, 0.4166667)
        self.assertEqual(s001_ad1.lattice.lengths_and_angles,
                                 s001.lattice.lengths_and_angles)
       # 2. test site_a = xyz input
        s001_ad2 = Slab.adsorb_atom(structure_a=s001, site_a=[1.5, 1.5, 3], atom= ['O'],
                                    distance=2, xyz=1)
        self.assertEqual(len(s001_ad2), 9)
        for i in xrange(len(s001_ad2)):
            if str(s001_ad2[i].specie) == 'O':
                print s001_ad2[i].frac_coords
                self.assertAlmostEqual(s001_ad2[i].a, 0.5)
                self.assertAlmostEqual(s001_ad2[i].b, 0.5)
                self.assertAlmostEqual(s001_ad2[i].c, 0.4166667)

if __name__ == "__main__":
    unittest.main()

<|endoftext|>


<commit_before>#
# This file is part of the PyMeasure package.
#
# Copyright (c) 2013-2021 PyMeasure Developers
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import logging

from logging import Handler

from .Qt import QtCore

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

class LogHandler(Handler):
    class Emitter(QtCore.QObject):
        record = QtCore.QSignal(object)

    def __init__(self):
        super().__init__()
        self.emitter = self.Emitter()

    def connect(self, *args, **kwargs):
        return self.emitter.record.connect(*args, **kwargs)

    def emit(self, record):
        self.emitter.record.emit(self.format(record))
<commit_msg>Explained purpose of class Emitter<commit_after>#
# This file is part of the PyMeasure package.
#
# Copyright (c) 2013-2021 PyMeasure Developers
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import logging

from logging import Handler

from .Qt import QtCore

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())

class LogHandler(Handler):
    # Class Emitter is added to keep compatibility with PySide2
    # 1. Signal needs to be class attribute of a QObject subclass
    # 2. logging Handler emit method clashes with QObject emit method
    # 3. As a consequence, the LogHandler cannot inherit both from
    # Handler and QObject
    # 4. A new utility class Emitter subclass of QObject is
    # introduced to handle record Signal and workaround the problem
    class Emitter(QtCore.QObject):
        record = QtCore.QSignal(object)

    def __init__(self):
        super().__init__()
        self.emitter = self.Emitter()

    def connect(self, *args, **kwargs):
        return self.emitter.record.connect(*args, **kwargs)

    def emit(self, record):
        self.emitter.record.emit(self.format(record))
<|endoftext|>


<commit_before>from .bot import Bot
from .utils import ToJsonMixin

from .buttons import *

@attr.s
class Element(ToJsonMixin):
    title = attr.ib()
    item_url = attr.ib(default=None)
    image_url = attr.ib(default=None)
    subtitle = attr.ib(default=None)
    buttons = attr.ib(default=None)
<commit_msg>Add QuickReply<commit_after>from .bot import Bot
from .utils import ToJsonMixin

from .buttons import *

@attr.s
class Element(ToJsonMixin):
    title = attr.ib()
    item_url = attr.ib(default=None)
    image_url = attr.ib(default=None)
    subtitle = attr.ib(default=None)
    buttons = attr.ib(default=None)

@attr.s
class QuickReply(ToJsonMixin):
    """
    See https://developers.facebook.com/docs/messenger-platform/send-api-reference/quick-replies

    You may not give the payload and it'll be set to your title automatically.
    """
    content_type = attr.ib()
    title = attr.ib(default=None)
    payload = attr.ib(default=None)
    image_url = attr.ib(default=None)

    def __attrs_post_init__(self):
        assert self.content_type in {'text', 'location'}
        assert self.content_type == 'location' or self.title

        if not self.payload:
            self.payload = self.title
<|endoftext|>


<commit_before>import logging
import re
from copy import deepcopy

from .evaluate import ConditionBlock

log = logging.getLogger(__name__)

class TraceableDict(dict):

    """
    Python dict-based class that tracks any changes.
    Format:
        {"action": "add", "key": <key>, "value": <new-value>},
        {"action": "remove", "key": <key>, "value": <old-value>},
        {"action": "update", "key": <key>, "old_value": <old-value>,
                                           "new_value": <new-value>},
    """

    def __init__(self, dict2):
        super().__init__(deepcopy(dict2))
        self._changes = []

    def __setitem__(self, key, value):
        if key not in self:
            self._changes.append({
                'action': 'add',
                'key': key,
                'value': deepcopy(value)
            })
        elif self[key] != value:
            self._changes.append({
                'action': 'update',
                'key': key,
                'old_value': deepcopy(self[key]),
                'new_value': deepcopy(value)
            })
        super().__setitem__(key, value)

    def __delitem__(self, key):
        self._changes.append({
            'action': 'remove',
            'key': key,
            'value': deepcopy(self[key])
        })
        super().__delitem__(key)

    def update(self, dict2=None, **kwargs):
        for key in dict2:
            if key not in self:
                self._changes.append({
                    'action': 'add',
                    'key': key,
                    'value': deepcopy(dict2[key])
                })
            elif self[key] != dict2[key]:
                self._changes.append({
                    'action': 'update',
                    'key': key,
                    'old_value': deepcopy(self[key]),
                    'new_value': deepcopy(dict2[key])
                })
        super().update(dict2, **kwargs)

    @property
    def changes(self):
        return self._changes

# Inspired from https://github.com/faif/python-patterns/blob/master/registry.py
class _RegisteredRule(type):
    """
    This metaclass is designed to register automatically public rule classes
    that inherit from `_Rules()`.
    """
    _registry = {}

    def __new__(mcs, name, bases, attrs):
        cls = type.__new__(mcs, name, bases, attrs)
        name = cls.TYPENAME or cls.__name__.lower()
        if not name.startswith('_'):
            cls.TYPENAME = name
            mcs._registry[name] = cls
        return cls

    @classmethod
    def get(mcs, rtype):
        return mcs._registry[rtype]

class Converter(object):

    """
    A sequence of `_Rule` objects intended to be applied on a dict.
    """

    def __init__(self, rules=None):
        self._rules = rules or []

    @property
    def rules(self):
        return self._rules

    @classmethod
    def from_dict(cls, config):
        """
        Create a `Converter` object from dict 'config'. The dict must look like
        the following:
        {
            "rules": [
                {"type": <rule-type-name>, "fieldname": <name>, ...},
                {"type": <rule-type-name>, "fieldname": <name>, ...},
                {"type": "condition-block", "conditions": [...]},
                {"type": <rule-type-name>, "fieldname": <name>, ...},
            ]
        }
        """
        rules = []
        for rule in config['rules']:
            params = rule.copy()
            rtype = params.pop('type')
            rule_cls = _RegisteredRule.get(rtype)
            rules.append(rule_cls(**params))
        return cls(rules=rules)

    def apply(self, data):
        rules = []
        errors = False
        for rule in self.rules:
            diff = rule.apply(data)
            rules.append(diff)
            if 'error' in diff:
                errors = True

        return {'rules': rules, 'errors': errors}

class FactoryConditionBlock(ConditionBlock, metaclass=_RegisteredRule):

    TYPENAME = 'condition-block'

    def __init__(self, conditions):
        super().__init__(conditions)
        self._changes = {'type': self.TYPENAME, 'conditions': []}

    def condition_validated(self, rules, data):
        """
        Apply rules on data upon validating a condition.
        """
        diff = Converter.from_dict({'rules': rules}).apply(data)
        self._changes['conditions'] = diff['rules']

    def apply(self, data):
        super().apply(data)
        return self._changes

class _Rule(metaclass=_RegisteredRule):

    """
    A rule extracts an entry from a dict , performs an operation on the dict
    and updates the dict with the result of that operation (in-place update).
    Only one field from the input dict can be processed within a rule.
    """

    # Public subclasses of `_Rule` will have a TYPENAME attr set to classname
    # (lowercase) if not set in the subclass itself.
    TYPENAME = None

    def __init__(self, fieldname, *args, **kwargs):
        self.fieldname = fieldname
        self._configure(*args, **kwargs)

    def _configure(self, *args, **kwargs):
        """
        Configure the rule to be applied (e.g. compile a regexp) so that
        everything is ready at runtime to apply it.
        """
        raise NotImplementedError

    @staticmethod
    def track_changes(func):
        """
        Decorator for `Rule.apply(<data>)` methods that return a JSON-formatted
        diff of the changes made by the method itself.
        """
        def wrapper(self, data):
            # Handle data through a traceable dict
            tracker = TraceableDict(data)
            try:
                func(self, tracker)
            except RegexpRuleError as exc:
                return {
                    'type': self.TYPENAME, 'changes': tracker.changes,
                    'error': 'regexp_rule_error', 'error_details': str(exc)
                }
            # We want to keep the object reference
            data.clear()
            data.update(tracker)
            return {'type': self.TYPENAME, 'changes': tracker.changes}

        return wrapper

    def apply(self, data):
        """
        Execute an operation on one field of the dict `data` and returns an
        diff.
        """
        raise NotImplementedError

class _RegexpRule(_Rule):

    """
    Define a regular expression that shall be applied through several kinds of
    operations to a string.
    If the regexp pattern has named capturing parenthesis it is used to update
    the `data` dict passed to `apply()`. If a group from the match object and
    a key from `data` have the same name, the group overrides the existing
    dict value.
    """

    def _configure(self, pattern, flags=0):
        self.regexp = re.compile(pattern, flags=flags)

    def _run_regexp(self, string):
        raise NotImplementedError

    @_Rule.track_changes
    def apply(self, data):
        try:
            string = data[self.fieldname]
        except KeyError:
            # No data to process
            log.debug('Regexp : unknown field %s, ignoring', self.fieldname)
        else:
            resdict = self._run_regexp(string)
            data.update(resdict)

class RegexpRuleError(TypeError):
    pass

class Extract(_RegexpRule):

    """
    Scan through a string looking for the first location where the regular
    expression pattern produces a match. What matters is the resulting dict of
    captured substrings.
    """

    def _configure(self, pattern, flags=0, pos=None, endpos=None):
        super()._configure(pattern, flags=flags)
        self._pos_args = tuple(f for f in (pos, endpos) if f is not None)

    def _run_regexp(self, string):
        if not self.regexp.groupindex:
            raise RegexpRuleError("regex is invalid, ensure a group is captured")
        args = (string,) + self._pos_args
        match = self.regexp.search(*args)
        if match is not None:
            return match.groupdict()
        else:
            return {}

class Sub(_RegexpRule):

    """
    Build a string obtained by replacing the leftmost non-overlapping
    occurrences of regexp pattern in fieldname from `data` by a replacement
    string.
    """

    def _configure(self, pattern, repl, flags=0, count=0):
        super()._configure(pattern, flags=flags)
        self.repl, self.count = repl, count

    def _run_regexp(self, string):
        res = self.regexp.sub(self.repl, string, count=self.count)
        return {self.fieldname: res}

class Set(_Rule):

    """
    Set or update a field in the `data` dict.
    """

    def _configure(self, value):
        self.value = value

    @_Rule.track_changes
    def apply(self, data):
        data[self.fieldname] = self.value

class Copy(_Rule):

    """
    Copy a data into another field.
    """

    def _configure(self, copy):
        self.copy = copy

    @_Rule.track_changes
    def apply(self, data):
        try:
            data[self.copy] = data[self.fieldname]
        except KeyError:
            log.debug('Copy : unknown field %s, ignoring', self.fieldname)

class Unset(_Rule):

    """
    Remove a field from the `data` dict.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        try:
            del data[self.fieldname]
        except KeyError:
            log.debug('Unset : unknown field %s, ignoring', self.fieldname)

class Lookup(_Rule):

    """
    Implements a dead-simple lookup table which perform case sensitive
    (default) or insensitive (icase=True) lookups.
    """

    def _configure(self, table=None, icase=False):
        self.table = table or dict()
        self.icase = icase
        if icase:
            self.table = {k.lower(): v for k, v in table.items()}

    @_Rule.track_changes
    def apply(self, data):
        """
        The 1st entry in the lookup table that matches the string from `data`
        replaces it. Note that entries in the lookup table are evaluated in an
        unpredictable order.
        """
        try:
            fieldval = str(data[self.fieldname])
            if self.icase:
                fieldval = fieldval.lower()
            data[self.fieldname] = self.table[fieldval]
        except KeyError as err:
            log.debug("Lookup: fieldname '%s' not in data, ignoring", err)
            return

class Lower(_Rule):

    """
    Lower case a string.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        fieldval = data[self.fieldname]
        data[self.fieldname] = fieldval.lower()

class Upper(_Rule):

    """
    Upper case a string.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        fieldval = data[self.fieldname]
        data[self.fieldname] = fieldval.upper()
<commit_msg>Failsafe in case the rule does not return any diff<commit_after>import logging
import re
from copy import deepcopy

from .evaluate import ConditionBlock

log = logging.getLogger(__name__)

class TraceableDict(dict):

    """
    Python dict-based class that tracks any changes.
    Format:
        {"action": "add", "key": <key>, "value": <new-value>},
        {"action": "remove", "key": <key>, "value": <old-value>},
        {"action": "update", "key": <key>, "old_value": <old-value>,
                                           "new_value": <new-value>},
    """

    def __init__(self, dict2):
        super().__init__(deepcopy(dict2))
        self._changes = []

    def __setitem__(self, key, value):
        if key not in self:
            self._changes.append({
                'action': 'add',
                'key': key,
                'value': deepcopy(value)
            })
        elif self[key] != value:
            self._changes.append({
                'action': 'update',
                'key': key,
                'old_value': deepcopy(self[key]),
                'new_value': deepcopy(value)
            })
        super().__setitem__(key, value)

    def __delitem__(self, key):
        self._changes.append({
            'action': 'remove',
            'key': key,
            'value': deepcopy(self[key])
        })
        super().__delitem__(key)

    def update(self, dict2=None, **kwargs):
        for key in dict2:
            if key not in self:
                self._changes.append({
                    'action': 'add',
                    'key': key,
                    'value': deepcopy(dict2[key])
                })
            elif self[key] != dict2[key]:
                self._changes.append({
                    'action': 'update',
                    'key': key,
                    'old_value': deepcopy(self[key]),
                    'new_value': deepcopy(dict2[key])
                })
        super().update(dict2, **kwargs)

    @property
    def changes(self):
        return self._changes

# Inspired from https://github.com/faif/python-patterns/blob/master/registry.py
class _RegisteredRule(type):
    """
    This metaclass is designed to register automatically public rule classes
    that inherit from `_Rules()`.
    """
    _registry = {}

    def __new__(mcs, name, bases, attrs):
        cls = type.__new__(mcs, name, bases, attrs)
        name = cls.TYPENAME or cls.__name__.lower()
        if not name.startswith('_'):
            cls.TYPENAME = name
            mcs._registry[name] = cls
        return cls

    @classmethod
    def get(mcs, rtype):
        return mcs._registry[rtype]

class Converter(object):

    """
    A sequence of `_Rule` objects intended to be applied on a dict.
    """

    def __init__(self, rules=None):
        self._rules = rules or []

    @property
    def rules(self):
        return self._rules

    @classmethod
    def from_dict(cls, config):
        """
        Create a `Converter` object from dict 'config'. The dict must look like
        the following:
        {
            "rules": [
                {"type": <rule-type-name>, "fieldname": <name>, ...},
                {"type": <rule-type-name>, "fieldname": <name>, ...},
                {"type": "condition-block", "conditions": [...]},
                {"type": <rule-type-name>, "fieldname": <name>, ...},
            ]
        }
        """
        rules = []
        for rule in config['rules']:
            params = rule.copy()
            rtype = params.pop('type')
            rule_cls = _RegisteredRule.get(rtype)
            rules.append(rule_cls(**params))
        return cls(rules=rules)

    def apply(self, data):
        rules = []
        errors = False
        for rule in self.rules:
            diff = rule.apply(data)
            rules.append(diff)
            if diff is not None and 'error' in diff:
                errors = True

        return {'rules': rules, 'errors': errors}

class FactoryConditionBlock(ConditionBlock, metaclass=_RegisteredRule):

    TYPENAME = 'condition-block'

    def __init__(self, conditions):
        super().__init__(conditions)
        self._changes = {'type': self.TYPENAME, 'conditions': []}

    def condition_validated(self, rules, data):
        """
        Apply rules on data upon validating a condition.
        """
        diff = Converter.from_dict({'rules': rules}).apply(data)
        self._changes['conditions'] = diff['rules']

    def apply(self, data):
        super().apply(data)
        return self._changes

class _Rule(metaclass=_RegisteredRule):

    """
    A rule extracts an entry from a dict , performs an operation on the dict
    and updates the dict with the result of that operation (in-place update).
    Only one field from the input dict can be processed within a rule.
    """

    # Public subclasses of `_Rule` will have a TYPENAME attr set to classname
    # (lowercase) if not set in the subclass itself.
    TYPENAME = None

    def __init__(self, fieldname, *args, **kwargs):
        self.fieldname = fieldname
        self._configure(*args, **kwargs)

    def _configure(self, *args, **kwargs):
        """
        Configure the rule to be applied (e.g. compile a regexp) so that
        everything is ready at runtime to apply it.
        """
        raise NotImplementedError

    @staticmethod
    def track_changes(func):
        """
        Decorator for `Rule.apply(<data>)` methods that return a JSON-formatted
        diff of the changes made by the method itself.
        """
        def wrapper(self, data):
            # Handle data through a traceable dict
            tracker = TraceableDict(data)
            try:
                func(self, tracker)
            except RegexpRuleError as exc:
                return {
                    'type': self.TYPENAME, 'changes': tracker.changes,
                    'error': 'regexp_rule_error', 'error_details': str(exc)
                }
            # We want to keep the object reference
            data.clear()
            data.update(tracker)
            return {'type': self.TYPENAME, 'changes': tracker.changes}

        return wrapper

    def apply(self, data):
        """
        Execute an operation on one field of the dict `data` and returns an
        diff.
        """
        raise NotImplementedError

class _RegexpRule(_Rule):

    """
    Define a regular expression that shall be applied through several kinds of
    operations to a string.
    If the regexp pattern has named capturing parenthesis it is used to update
    the `data` dict passed to `apply()`. If a group from the match object and
    a key from `data` have the same name, the group overrides the existing
    dict value.
    """

    def _configure(self, pattern, flags=0):
        self.regexp = re.compile(pattern, flags=flags)

    def _run_regexp(self, string):
        raise NotImplementedError

    @_Rule.track_changes
    def apply(self, data):
        try:
            string = data[self.fieldname]
        except KeyError:
            # No data to process
            log.debug('Regexp : unknown field %s, ignoring', self.fieldname)
        else:
            resdict = self._run_regexp(string)
            data.update(resdict)

class RegexpRuleError(TypeError):
    pass

class Extract(_RegexpRule):

    """
    Scan through a string looking for the first location where the regular
    expression pattern produces a match. What matters is the resulting dict of
    captured substrings.
    """

    def _configure(self, pattern, flags=0, pos=None, endpos=None):
        super()._configure(pattern, flags=flags)
        self._pos_args = tuple(f for f in (pos, endpos) if f is not None)

    def _run_regexp(self, string):
        if not self.regexp.groupindex:
            raise RegexpRuleError("regex is invalid, ensure a group is captured")
        args = (string,) + self._pos_args
        match = self.regexp.search(*args)
        if match is not None:
            return match.groupdict()
        else:
            return {}

class Sub(_RegexpRule):

    """
    Build a string obtained by replacing the leftmost non-overlapping
    occurrences of regexp pattern in fieldname from `data` by a replacement
    string.
    """

    def _configure(self, pattern, repl, flags=0, count=0):
        super()._configure(pattern, flags=flags)
        self.repl, self.count = repl, count

    def _run_regexp(self, string):
        res = self.regexp.sub(self.repl, string, count=self.count)
        return {self.fieldname: res}

class Set(_Rule):

    """
    Set or update a field in the `data` dict.
    """

    def _configure(self, value):
        self.value = value

    @_Rule.track_changes
    def apply(self, data):
        data[self.fieldname] = self.value

class Copy(_Rule):

    """
    Copy a data into another field.
    """

    def _configure(self, copy):
        self.copy = copy

    @_Rule.track_changes
    def apply(self, data):
        try:
            data[self.copy] = data[self.fieldname]
        except KeyError:
            log.debug('Copy : unknown field %s, ignoring', self.fieldname)

class Unset(_Rule):

    """
    Remove a field from the `data` dict.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        try:
            del data[self.fieldname]
        except KeyError:
            log.debug('Unset : unknown field %s, ignoring', self.fieldname)

class Lookup(_Rule):

    """
    Implements a dead-simple lookup table which perform case sensitive
    (default) or insensitive (icase=True) lookups.
    """

    def _configure(self, table=None, icase=False):
        self.table = table or dict()
        self.icase = icase
        if icase:
            self.table = {k.lower(): v for k, v in table.items()}

    @_Rule.track_changes
    def apply(self, data):
        """
        The 1st entry in the lookup table that matches the string from `data`
        replaces it. Note that entries in the lookup table are evaluated in an
        unpredictable order.
        """
        try:
            fieldval = str(data[self.fieldname])
            if self.icase:
                fieldval = fieldval.lower()
            data[self.fieldname] = self.table[fieldval]
        except KeyError as err:
            log.debug("Lookup: fieldname '%s' not in data, ignoring", err)
            return

class Lower(_Rule):

    """
    Lower case a string.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        fieldval = data[self.fieldname]
        data[self.fieldname] = fieldval.lower()

class Upper(_Rule):

    """
    Upper case a string.
    """

    def _configure(self):
        pass

    @_Rule.track_changes
    def apply(self, data):
        fieldval = data[self.fieldname]
        data[self.fieldname] = fieldval.upper()
<|endoftext|>