<commit_before>import os
import abc
import boto3
import shutil
import logging
import tarfile
import zipfile
import urllib2
import tempfile
import subprocess

from botocore.client import ClientError
from watchmaker.exceptions import SystemFatal as exceptionhandler

class ManagerBase(object):
    """
    Base class for operating system managers.  This forces all child classes to require consistent methods
    for coherence.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def _get_s3_file(self, url, bucket_name, key_name, destination):
        """

        :param url:
        :param bucket_name:
        :param key_name:
        :param destination:
        :return:
        """

        return

    @abc.abstractmethod
    def download_file(self, url, filename, sourceiss3bucket):
        """

        :param url:
        :param filename:
        :param sourceiss3bucket:
        :return:
        """

        return

    @abc.abstractmethod
    def create_working_dir(self, basedir, prefix):
        """

        :param basedir:
        :param prefix:
        :return:
        """

        return

    @abc.abstractmethod
    def call_process(self, cmd):
        return

    @abc.abstractmethod
    def cleanup(self):
        """

        :return:
        """

        return

    @abc.abstractmethod
    def extract_contents(self, filepath, to_directory, create_dir):
        """

        :param filepath:
        :param to_directory:
        :param create_dir:
        :return:
        """

        return

class LinuxManager(ManagerBase):
    """
    This is the base import for Linux Managers.  It serves as a foundational class to keep OS consitency.

    """

    def __init__(self):
        self.workingdir = None

    def _get_s3_file(self, url, bucket_name, key_name, destination):

        try:
            s3 = boto3.resource("s3")
            s3.meta.client.head_bucket(Bucket=bucket_name)
            s3.Object(bucket_name, key_name).download_file(destination)
        except ClientError as exc:
            logging.error('Bucket does not exist.\n'
                          'bucket = {0}\n'
                          'Exception: {1}'
                          .format(bucket_name, exc))
            raise SystemError('Bucket does not exist.\n'
                              'bucket = {0}\n'
                              'Exception: {1}'
                              .format(bucket_name, exc))
        except Exception as exc:
            logging.error('Unable to download file from S3 bucket.\n'
                          'url = {0}\n'
                          'bucket = {1}\n'
                          'key = {2}\n'
                          'file = {3}\n'
                          'Exception: {4}'
                          .format(url, bucket_name, key_name,
                                  destination, exc))
            raise SystemError('Unable to download file from S3 bucket.\n'
                              'url = {0}\n'
                              'bucket = {1}\n'
                              'key = {2}\n'
                              'file = {3}\n'
                              'Exception: {4}'
                              .format(url, bucket_name, key_name,
                                      destination, exc))

    def _install_from_yum(self, packages):
        """

        :param packages:
        :return:
        """

        yum_cmd = ['sudo', 'yum', '-y', 'install']
        if isinstance(packages, list):
            yum_cmd.extend(packages)
        else:
            yum_cmd.append(packages)
        rsp = subprocess.call(yum_cmd)
        logging.debug(packages)
        logging.debug('Return code of yum install: {0}'.format(rsp))

        if rsp != 0:
            exceptionhandler('Installing Salt from Yum has failed!')

    def download_file(self, url, filename, sourceiss3bucket=False):
        """

        :param url:
        :param filename:
        :param sourceiss3bucket:
        :return:
        """
        conn = None

        logging.debug('Downloading: {0}'.format(url))
        logging.debug('Destination: {0}'.format(filename))
        logging.debug('S3: {0}'.format(sourceiss3bucket))

        # TODO Rework this to properly reflect logic flow cleanly.
        if sourceiss3bucket:
            bucket_name = url.split('/')[3]
            key_name = '/'.join(url.split('/')[4:])

            logging.debug('Bucket Name: {0}'.format(bucket_name))
            logging.debug('key_name: {0}'.format(key_name))

            try:
                s3 = boto3.resource("s3")
                s3.meta.client.head_bucket(Bucket=bucket_name)
                s3.Object(bucket_name, key_name).download_file(filename)
            except (NameError, ClientError):
                logging.error('NameError: {0}'.format(ClientError))
                try:
                    bucket_name = url.split('/')[2].split('.')[0]
                    key_name = '/'.join(url.split('/')[3:])
                    s3 = boto3.resource("s3")
                    s3.meta.client.head_bucket(Bucket=bucket_name)
                    s3.Object(bucket_name, key_name).download_file(filename)
                except Exception as exc:
                    logging.error('Unable to download file from S3 bucket.\n'
                                  'url = {0}\n'
                                  'bucket = {1}\n'
                                  'key = {2}\n'
                                  'file = {3}\n'
                                  'Exception: {4}'
                                  .format(url, bucket_name, key_name, filename, exc))
                    raise SystemError('Unable to download file from S3 bucket.\n'
                                      'url = {0}\n'
                                      'bucket = {1}\n'
                                      'key = {2}\n'
                                      'file = {3}\n'
                                      'Exception: {4}'
                                      .format(url, bucket_name, key_name, filename, exc))
            except Exception as exc:
                logging.error('Unable to download file from S3 bucket.\n'
                              'url = {0}\n'
                              'bucket = {1}\n'
                              'key = {2}\n'
                              'file = {3}\n'
                              'Exception: {4}'
                              .format(url, bucket_name, key_name, filename, exc))
                raise SystemError('Unable to download file from S3 bucket.\n'
                                  'url = {0}\n'
                                  'bucket = {1}\n'
                                  'key = {2}\n'
                                  'file = {3}\n'
                                  'Exception: {4}'
                                  .format(url, bucket_name, key_name, filename, exc))
            logging.debug('Downloaded file from S3 bucket -- \n'
                          '    url      = {0}\n'
                          '    filename = {1}'.format(url, filename))
        else:
            try:
                response = urllib2.urlopen(url)
                with open(filename, 'wb') as outfile:
                    shutil.copyfileobj(response, outfile)
            except Exception as exc:
                # TODO: Update `except` logic
                logging.error('Unable to download file from web server.\n'
                              'url = {0}\n'
                              'filename = {1}\n'
                              'Exception: {2}'
                              .format(url, filename, exc))
                raise SystemError('Unable to download file from web server.\n'
                                  'url = {0}\n'
                                  'filename = {1}\n'
                                  'Exception: {2}'
                                  .format(url, filename, exc))
            logging.debug('Downloaded file from web server -- \n'
                          '    url      = {0}\n'
                          '    filename = {1}'.format(url, filename))

    def call_process(self, cmd):
        if not isinstance(cmd, list):
            exceptionhandler('Command is not a list.\n{0}'.format(str(cmd)))
        rsp = subprocess.call(cmd)

        if rsp != 0:
            exceptionhandler('Command failed.\n{0}'.format(str(cmd)))

    def create_working_dir(self, basedir, prefix):
        """
        Creates a directory in `basedir` with a prefix of `dirprefix`.
        The directory will have a random 5 character string appended to `dirprefix`.
        Returns the path to the working directory.

        :param prefix:
        :rtype : str
        :param basedir: str, the directory in which to create the working directory
        """

        logging.info('Creating a working directory.')
        workingdir = None
        original_umask = os.umask(0)
        try:
            workingdir = tempfile.mkdtemp(prefix=prefix, dir=basedir)
        except Exception as exc:
            exceptionhandler('Could not create workingdir in {0}.\n'
                             'Exception: {1}'.format(basedir, exc))
        logging.debug('Working directory: {0}'.format(workingdir))
        self.workingdir = workingdir
        os.umask(original_umask)

    def cleanup(self):
        """
        Removes temporary files loaded to the system.
            :return: bool
        """
        logging.info('+-' * 40)
        logging.info('Cleanup Time...')
        try:
            logging.debug('{0} being cleaned up.'.format(self.workingdir))
            shutil.rmtree(self.workingdir)
        except Exception as exc:
            # TODO: Update `except` logic
            logging.fatal('Cleanup Failed!\nException: {0}'.format(exc))
            exceptionhandler('Cleanup Failed.\nAborting.')

        logging.info('Removed temporary data in working directory -- {0}'.format(self.workingdir))
        logging.info('Exiting cleanup routine...')
        logging.info('-+' * 40)

    def extract_contents(self, filepath, to_directory='.', create_dir=None):
        """
        Extracts a compressed file to the specified directory.
        Supports files that end in .zip, .tar.gz, .tgz, tar.bz2, or tbz.
        :param create_dir:
        :param filepath: str, path to the compressed file
        :param to_directory: str, path to the target directory
        :raise ValueError: error raised if file extension is not supported
        """
        opener = None
        mode = None

        if filepath.endswith('.zip'):
            logging.debug('File Type: zip')
            opener, mode = zipfile.ZipFile, 'r'
        elif filepath.endswith('.tar.gz') or filepath.endswith('.tgz'):
            logging.debug('File Type: GZip Tar')
            opener, mode = tarfile.open, 'r:gz'
        elif filepath.endswith('.tar.bz2') or filepath.endswith('.tbz'):
            logging.debug('File Type: Bzip Tar')
            opener, mode = tarfile.open, 'r:bz2'
        else:
            exceptionhandler('{0}'.format('Could not extract `"{0}`" as no appropriate '
                             'extractor is found'.format(filepath)))

        if create_dir:
            to_directory = os.sep.join((to_directory,
                                        '.'.join(filepath.split(os.sep)[-1].split('.')[:-1])))

        try:
            os.makedirs(to_directory)
        except OSError:
            if not os.path.isdir(to_directory):
                raise

        cwd = os.getcwd()
        os.chdir(to_directory)

        try:
            openfile = opener(filepath, mode)
            try:
                openfile.extractall()
            finally:
                openfile.close()
        finally:
            os.chdir(cwd)

        print('Extracted file -- \n'
              '    source = {0}\n'
              '    dest   = {1}'.format(filepath, to_directory))

class WindowsManager(ManagerBase):
    """

    """

    def __init__(self):
        pass

    def _get_s3_file(self, url, bucket_name, key_name, destination):
        pass

    def download_file(self, url, filename, sourceiss3bucket):
        pass

    def call_process(self, cmd):
        pass

    def create_working_dir(self, basedir, prefix):
        pass

    def cleanup(self):
        pass

    def extract_contents(self, filepath, to_directory, create_dir):
        pass

class WorkersManagerBase(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def _worker_execution(self):
        return

    @abc.abstractmethod
    def _worker_validation(self):
        return

    @abc.abstractmethod
    def worker_cadence(self):
        return

    @abc.abstractmethod
    def cleanup(self):
        return
<commit_msg>Fixing issues after merge.<commit_after>import os
import abc
import boto3
import shutil
import logging
import tarfile
import zipfile
import urllib2
import tempfile
import subprocess

from botocore.client import ClientError
from watchmaker.exceptions import SystemFatal as exceptionhandler

class ManagerBase(object):
    """
    Base class for operating system managers.  This forces all child classes to require consistent methods
    for coherence.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def __init__(self):
        return

    @staticmethod
    @abc.abstractmethod
    def _get_s3_file(self, url, bucket_name, key_name, destination):
        """

        :param url:
        :param bucket_name:
        :param key_name:
        :param destination:
        :return:
        """

        try:
            s3 = boto3.resource("s3")
            s3.meta.client.head_bucket(Bucket=bucket_name)
            s3.Object(bucket_name, key_name).download_file(destination)
        except ClientError as exc:
            logging.error('Bucket does not exist.\n'
                          'bucket = {0}\n'
                          'Exception: {1}'
                          .format(bucket_name, exc))
            raise SystemError('Bucket does not exist.\n'
                              'bucket = {0}\n'
                              'Exception: {1}'
                              .format(bucket_name, exc))
        except Exception as exc:
            logging.error('Unable to download file from S3 bucket.\n'
                          'url = {0}\n'
                          'bucket = {1}\n'
                          'key = {2}\n'
                          'file = {3}\n'
                          'Exception: {4}'
                          .format(url, bucket_name, key_name,
                                  destination, exc))
            raise SystemError('Unable to download file from S3 bucket.\n'
                              'url = {0}\n'
                              'bucket = {1}\n'
                              'key = {2}\n'
                              'file = {3}\n'
                              'Exception: {4}'
                              .format(url, bucket_name, key_name,
                                      destination, exc))

    @abc.abstractmethod
    def download_file(self, url, filename, sourceiss3bucket):
        """

        :param url:
        :param filename:
        :param sourceiss3bucket:
        :return:
        """

        return

    @abc.abstractmethod
    def create_working_dir(self, basedir, prefix):
        """

        :param basedir:
        :param prefix:
        :return:
        """

        return

    @abc.abstractmethod
    def call_process(self, cmd):
        return

    @abc.abstractmethod
    def cleanup(self):
        """

        :return:
        """

        return

    @abc.abstractmethod
    def extract_contents(self, filepath, to_directory, create_dir):
        """

        :param filepath:
        :param to_directory:
        :param create_dir:
        :return:
        """

        return

class LinuxManager(ManagerBase):
    """
    This is the base import for Linux Managers.  It serves as a foundational class to keep OS consitency.

    """

    def __init__(self):
        super(LinuxManager, self).__init__()
        self.workingdir = None

    def _install_from_yum(self, packages):
        """

        :param packages:
        :return:
        """

        yum_cmd = ['sudo', 'yum', '-y', 'install']
        if isinstance(packages, list):
            yum_cmd.extend(packages)
        else:
            yum_cmd.append(packages)
        rsp = subprocess.call(yum_cmd)
        logging.debug(packages)
        logging.debug('Return code of yum install: {0}'.format(rsp))

        if rsp != 0:
            exceptionhandler('Installing Salt from Yum has failed!')

    def download_file(self, url, filename, sourceiss3bucket=False):
        """

        :param url:
        :param filename:
        :param sourceiss3bucket:
        :return:
        """
        conn = None

        logging.debug('Downloading: {0}'.format(url))
        logging.debug('Destination: {0}'.format(filename))
        logging.debug('S3: {0}'.format(sourceiss3bucket))

        # TODO Rework this to properly reflect logic flow cleanly.
        if sourceiss3bucket:
            bucket_name = url.split('/')[3]
            key_name = '/'.join(url.split('/')[4:])

            logging.debug('Bucket Name: {0}'.format(bucket_name))
            logging.debug('key_name: {0}'.format(key_name))

            try:
                s3 = boto3.resource("s3")
                s3.meta.client.head_bucket(Bucket=bucket_name)
                s3.Object(bucket_name, key_name).download_file(filename)
            except (NameError, ClientError):
                logging.error('NameError: {0}'.format(ClientError))
                try:
                    bucket_name = url.split('/')[2].split('.')[0]
                    key_name = '/'.join(url.split('/')[3:])
                    s3 = boto3.resource("s3")
                    s3.meta.client.head_bucket(Bucket=bucket_name)
                    s3.Object(bucket_name, key_name).download_file(filename)
                except Exception as exc:
                    logging.error('Unable to download file from S3 bucket.\n'
                                  'url = {0}\n'
                                  'bucket = {1}\n'
                                  'key = {2}\n'
                                  'file = {3}\n'
                                  'Exception: {4}'
                                  .format(url, bucket_name, key_name, filename, exc))
                    raise SystemError('Unable to download file from S3 bucket.\n'
                                      'url = {0}\n'
                                      'bucket = {1}\n'
                                      'key = {2}\n'
                                      'file = {3}\n'
                                      'Exception: {4}'
                                      .format(url, bucket_name, key_name, filename, exc))
            except Exception as exc:
                logging.error('Unable to download file from S3 bucket.\n'
                              'url = {0}\n'
                              'bucket = {1}\n'
                              'key = {2}\n'
                              'file = {3}\n'
                              'Exception: {4}'
                              .format(url, bucket_name, key_name, filename, exc))
                raise SystemError('Unable to download file from S3 bucket.\n'
                                  'url = {0}\n'
                                  'bucket = {1}\n'
                                  'key = {2}\n'
                                  'file = {3}\n'
                                  'Exception: {4}'
                                  .format(url, bucket_name, key_name, filename, exc))
            logging.debug('Downloaded file from S3 bucket -- \n'
                          '    url      = {0}\n'
                          '    filename = {1}'.format(url, filename))
        else:
            try:
                response = urllib2.urlopen(url)
                with open(filename, 'wb') as outfile:
                    shutil.copyfileobj(response, outfile)
            except Exception as exc:
                # TODO: Update `except` logic
                logging.error('Unable to download file from web server.\n'
                              'url = {0}\n'
                              'filename = {1}\n'
                              'Exception: {2}'
                              .format(url, filename, exc))
                raise SystemError('Unable to download file from web server.\n'
                                  'url = {0}\n'
                                  'filename = {1}\n'
                                  'Exception: {2}'
                                  .format(url, filename, exc))
            logging.debug('Downloaded file from web server -- \n'
                          '    url      = {0}\n'
                          '    filename = {1}'.format(url, filename))

    def call_process(self, cmd):
        if not isinstance(cmd, list):
            exceptionhandler('Command is not a list.\n{0}'.format(str(cmd)))
        rsp = subprocess.call(cmd)

        if rsp != 0:
            exceptionhandler('Command failed.\n{0}'.format(str(cmd)))

    def create_working_dir(self, basedir, prefix):
        """
        Creates a directory in `basedir` with a prefix of `dirprefix`.
        The directory will have a random 5 character string appended to `dirprefix`.
        Returns the path to the working directory.

        :param prefix:
        :rtype : str
        :param basedir: str, the directory in which to create the working directory
        """

        logging.info('Creating a working directory.')
        workingdir = None
        original_umask = os.umask(0)
        try:
            workingdir = tempfile.mkdtemp(prefix=prefix, dir=basedir)
        except Exception as exc:
            exceptionhandler('Could not create workingdir in {0}.\n'
                             'Exception: {1}'.format(basedir, exc))
        logging.debug('Working directory: {0}'.format(workingdir))
        self.workingdir = workingdir
        os.umask(original_umask)

    def cleanup(self):
        """
        Removes temporary files loaded to the system.
            :return: bool
        """
        logging.info('+-' * 40)
        logging.info('Cleanup Time...')
        try:
            logging.debug('{0} being cleaned up.'.format(self.workingdir))
            shutil.rmtree(self.workingdir)
        except Exception as exc:
            # TODO: Update `except` logic
            logging.fatal('Cleanup Failed!\nException: {0}'.format(exc))
            exceptionhandler('Cleanup Failed.\nAborting.')

        logging.info('Removed temporary data in working directory -- {0}'.format(self.workingdir))
        logging.info('Exiting cleanup routine...')
        logging.info('-+' * 40)

    def extract_contents(self, filepath, to_directory='.', create_dir=None):
        """
        Extracts a compressed file to the specified directory.
        Supports files that end in .zip, .tar.gz, .tgz, tar.bz2, or tbz.
        :param create_dir:
        :param filepath: str, path to the compressed file
        :param to_directory: str, path to the target directory
        :raise ValueError: error raised if file extension is not supported
        """
        opener = None
        mode = None

        if filepath.endswith('.zip'):
            logging.debug('File Type: zip')
            opener, mode = zipfile.ZipFile, 'r'
        elif filepath.endswith('.tar.gz') or filepath.endswith('.tgz'):
            logging.debug('File Type: GZip Tar')
            opener, mode = tarfile.open, 'r:gz'
        elif filepath.endswith('.tar.bz2') or filepath.endswith('.tbz'):
            logging.debug('File Type: Bzip Tar')
            opener, mode = tarfile.open, 'r:bz2'
        else:
            exceptionhandler('{0}'.format('Could not extract `"{0}`" as no appropriate '
                             'extractor is found'.format(filepath)))

        if create_dir:
            to_directory = os.sep.join((to_directory,
                                        '.'.join(filepath.split(os.sep)[-1].split('.')[:-1])))

        try:
            os.makedirs(to_directory)
        except OSError:
            if not os.path.isdir(to_directory):
                raise

        cwd = os.getcwd()
        os.chdir(to_directory)

        try:
            openfile = opener(filepath, mode)
            try:
                openfile.extractall()
            finally:
                openfile.close()
        finally:
            os.chdir(cwd)

        print('Extracted file -- \n'
              '    source = {0}\n'
              '    dest   = {1}'.format(filepath, to_directory))

class WindowsManager(ManagerBase):
    """

    """

    def __init__(self):
        super(WindowsManager, self).__init__()
        self.workingdir = None

    def download_file(self, url, filename, sourceiss3bucket):
        pass

    def call_process(self, cmd):
        pass

    def create_working_dir(self, basedir, prefix):
        pass

    def cleanup(self):
        pass

    def extract_contents(self, filepath, to_directory, create_dir):
        pass

class WorkersManagerBase(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def _worker_execution(self):
        return

    @abc.abstractmethod
    def _worker_validation(self):
        return

    @abc.abstractmethod
    def worker_cadence(self):
        return

    @abc.abstractmethod
    def cleanup(self):
        return
<|endoftext|>


<commit_before>import logging
from typing import Optional, List, Set, Tuple

import claripy

from . import PageBase
from angr.storage.memory_object import SimMemoryObject
from .cooperation import MemoryObjectMixin

l = logging.getLogger(name=__name__)

class ListPage(MemoryObjectMixin, PageBase):
    def __init__(self, memory=None, content=None, sinkhole=None, **kwargs):
        super().__init__(**kwargs)

        self.content: List[Optional[SimMemoryObject]] = content
        if content is None:
            if memory is not None:
                self.content: List[Optional[SimMemoryObject]] = [None] * memory.page_size  # TODO: this isn't the best

        self.sinkhole: Optional[SimMemoryObject] = sinkhole

    def copy(self, memo):
        o = super().copy(memo)
        o.content = list(self.content)
        o.sinkhole = self.sinkhole
        return o

    def load(self, addr, size=None, endness=None, page_addr=None, memory=None, cooperate=False, **kwargs):
        result = []
        last_seen = ...  # ;)

        # loop over the loading range. accumulate a result for each byte, but collapse results from adjacent bytes
        # using the same memory object
        for subaddr in range(addr, addr+size):
            item = self.content[subaddr]
            if item is None:
                item = self.sinkhole
            if item is not last_seen:
                if last_seen is None:
                    self._fill(result, subaddr, page_addr, endness, memory, **kwargs)
                result.append((subaddr + page_addr, item))
                last_seen = item

        if last_seen is None:
            self._fill(result, addr + size, page_addr, endness, memory, **kwargs)

        if not cooperate:
            result = self._force_load_cooperation(result, size, endness, memory=memory, **kwargs)
        return result

    def _fill(self, result, addr, page_addr, endness, memory, **kwargs):
        """
        Small utility function for behavior which is duplicated in load

        mutates result to generate a new memory object and replace the last entry in it, which is None. Then, it will
        insert the new memory object into self.content.
        """
        global_end_addr = addr + page_addr
        global_start_addr = result[-1][0]
        size = global_end_addr - global_start_addr
        new_ast = self._default_value(global_start_addr, size, name='%s_%x' % (memory.id, global_start_addr), key=(self.category, global_start_addr), memory=memory, **kwargs)
        new_item = SimMemoryObject(new_ast, global_start_addr, endness=endness, byte_width=memory.state.arch.byte_width if memory is not None else 8)
        for subaddr in range(global_start_addr - page_addr, addr):
            self.content[subaddr] = new_item
        result[-1] = (global_start_addr, new_item)

    def store(self, addr, data, size=None, endness=None, memory=None, cooperate=False, **kwargs):
        if not cooperate:
            data = self._force_store_cooperation(addr, data, size, endness, memory=memory, **kwargs)

        if size == len(self.content) and addr == 0:
            self.sinkhole = data
            self.content = [None]*len(self.content)
        else:
            for subaddr in range(addr, addr + size):
                self.content[subaddr] = data

    def merge(self, others: List['ListPage'], merge_conditions, common_ancestor=None, page_addr: int=None,
              memory=None):

        changed_offsets = set()
        for other in others:
            changed_offsets |= self.changed_bytes(other, page_addr)

        all_pages: List['ListPage'] = [self] + others
        if merge_conditions is None:
            merge_conditions = [None] * len(all_pages)

        merged_to = None
        merged_objects = set()
        merged_offsets = set()
        for b in sorted(changed_offsets):
            if merged_to is not None and not b >= merged_to:
                l.info("merged_to = %d ... already merged byte 0x%x", merged_to, b)
                continue
            l.debug("... on byte 0x%x", b)

            memory_objects = []
            unconstrained_in = []

            # first get a list of all memory objects at that location, and
            # all memories that don't have those bytes
            for sm, fv in zip(all_pages, merge_conditions):
                if sm._contains(b, page_addr):
                    l.info("... present in %s", fv)
                    memory_objects.append((sm.content[b], fv))
                else:
                    l.info("... not present in %s", fv)
                    unconstrained_in.append((sm, fv))

            mos = set(mo for mo, _ in memory_objects)
            mo_bases = set(mo.base for mo, _ in memory_objects)
            mo_lengths = set(mo.length for mo, _ in memory_objects)

            if not unconstrained_in and not (mos - merged_objects):
                continue

            # first, optimize the case where we are dealing with the same-sized memory objects
            if len(mo_bases) == 1 and len(mo_lengths) == 1 and not unconstrained_in:
                our_mo = self.content[b]
                to_merge = [(mo.object, fv) for mo, fv in memory_objects]

                # Update `merged_to`
                mo_base = list(mo_bases)[0]
                merged_to = mo_base + list(mo_lengths)[0]
                mo_length = memory_objects[0][0].length

                merged_val = self._merge_values(to_merge, mo_length, memory=memory)
                if merged_val is None:
                    # merge_values() determines that we should not attempt to merge this value
                    continue

                # do the replacement
                # TODO: Implement in-place replacement instead of calling store()
                # new_object = self._replace_memory_object(our_mo, merged_val, page_addr, memory.page_size)

                self.store(b,
                           SimMemoryObject(merged_val, mo_base, endness='Iend_BE'),
                           size=mo_length - (page_addr + b - mo_base),
                           cooperate=True
                           )
                # merged_objects.add(new_object)
                # merged_objects.update(mos)
                merged_offsets.add(b)

            else:
                # get the size that we can merge easily. This is the minimum of
                # the size of all memory objects and unallocated spaces.
                min_size = min([mo.length - (b + page_addr - mo.base) for mo, _ in memory_objects])
                for um, _ in unconstrained_in:
                    for i in range(0, min_size):
                        if um._contains(b + i, page_addr):
                            min_size = i
                            break
                merged_to = b + min_size
                l.info("... determined minimum size of %d", min_size)

                # Now, we have the minimum size. We'll extract/create expressions of that
                # size and merge them
                extracted = [(mo.bytes_at(page_addr+b, min_size), fv) for mo, fv in memory_objects] if min_size != 0 else []
                created = [
                    (self._default_value(None, min_size, name="merge_uc_%s_%x" % (uc.id, b), memory=memory),
                     fv) for
                    uc, fv in unconstrained_in
                ]
                to_merge = extracted + created

                merged_val = self._merge_values(to_merge, min_size, memory=memory)
                if merged_val is None:
                    continue

                self.store(b,
                           SimMemoryObject(merged_val, page_addr+b, endness='Iend_BE'),
                           size=min_size,
                           endness='Iend_BE', cooperate=True
                           )  # do not convert endianness again
                merged_offsets.add(b)

        return merged_offsets

    def changed_bytes(self, other: 'ListPage', page_addr: int=None):

        candidates: Set[int] = set()
        for i in range(len(self.content)):
            if self._contains(i, page_addr):
                candidates.add(i)
        for i in range(len(other.content)):
            if other._contains(i, page_addr):
                candidates.add(i)

        byte_width = 8  # TODO: Introduce self.state if we want to use self.state.arch.byte_width
        differences: Set[int] = set()
        for c in candidates:
            if not self._contains(c, page_addr) and other._contains(c, page_addr):
                differences.add(c)
            elif self._contains(c, page_addr) and not other._contains(c, page_addr):
                differences.add(c)
            else:
                if self.content[c] is None:
                    self.content[c] = SimMemoryObject(self.sinkhole.bytes_at(page_addr+c, 1), page_addr + c,
                                                      byte_width=byte_width, endness='Iend_BE')
                if other.content[c] is None:
                    other.content[c] = SimMemoryObject(other.sinkhole.bytes_at(page_addr+c, 1), page_addr + c,
                                                       byte_width=byte_width, endness='Iend_BE')
                if self._contains(c, page_addr) and self.content[c] != other.content[c]:
                    # Try to see if the bytes are equal
                    self_byte = self.content[c].bytes_at(page_addr + c, 1)
                    other_byte = other.content[c].bytes_at(page_addr + c, 1)
                    if self_byte is not other_byte:
                        #l.debug("%s: offset %x, two different bytes %s %s from %s %s", self.id, c,
                        #        self_byte, other_byte,
                        #        self[c].object.model, other[c].object.model)
                        differences.add(c)
                else:
                    # this means the byte is in neither memory
                    pass

        return differences

    def _contains(self, off: int, page_addr: int):
        if off >= len(self.content):
            return False
        if self.content[off] is not None:
            return True
        if self.sinkhole is None:
            return False
        return self.sinkhole.includes(page_addr + off)

    def _replace_mo(self, old_mo: SimMemoryObject, new_mo: SimMemoryObject, page_addr: int,
                    page_size: int) -> SimMemoryObject:
        if self.sinkhole is old_mo:
            self.sinkhole = new_mo
        else:
            start, end = self._resolve_range(old_mo, page_addr, page_size)
            for i in range(start, end):
                if self.content[i-page_addr] is old_mo:
                    self.content[i-page_addr] = new_mo
        return new_mo

    @staticmethod
    def _resolve_range(mo: SimMemoryObject, page_addr: int, page_size) -> Tuple[int,int]:
        start = max(mo.base, page_addr)
        end = min(mo.last_addr + 1, page_addr + page_size)
        if end <= start:
            l.warning("Nothing left of the memory object to store in SimPage.")
        return start, end

    def _get_object(self, start: int, page_addr: int) -> Optional[SimMemoryObject]:
        mo = self.content[start]
        if mo is None:
            return None
        if mo.includes(start + page_addr):
            return mo
        return None
<commit_msg>Correct the value of merged_to in ListPage.merge() (#2340)<commit_after>import logging
from typing import Optional, List, Set, Tuple

import claripy

from . import PageBase
from angr.storage.memory_object import SimMemoryObject
from .cooperation import MemoryObjectMixin

l = logging.getLogger(name=__name__)

class ListPage(MemoryObjectMixin, PageBase):
    def __init__(self, memory=None, content=None, sinkhole=None, **kwargs):
        super().__init__(**kwargs)

        self.content: List[Optional[SimMemoryObject]] = content
        if content is None:
            if memory is not None:
                self.content: List[Optional[SimMemoryObject]] = [None] * memory.page_size  # TODO: this isn't the best

        self.sinkhole: Optional[SimMemoryObject] = sinkhole

    def copy(self, memo):
        o = super().copy(memo)
        o.content = list(self.content)
        o.sinkhole = self.sinkhole
        return o

    def load(self, addr, size=None, endness=None, page_addr=None, memory=None, cooperate=False, **kwargs):
        result = []
        last_seen = ...  # ;)

        # loop over the loading range. accumulate a result for each byte, but collapse results from adjacent bytes
        # using the same memory object
        for subaddr in range(addr, addr+size):
            item = self.content[subaddr]
            if item is None:
                item = self.sinkhole
            if item is not last_seen:
                if last_seen is None:
                    self._fill(result, subaddr, page_addr, endness, memory, **kwargs)
                result.append((subaddr + page_addr, item))
                last_seen = item

        if last_seen is None:
            self._fill(result, addr + size, page_addr, endness, memory, **kwargs)

        if not cooperate:
            result = self._force_load_cooperation(result, size, endness, memory=memory, **kwargs)
        return result

    def _fill(self, result, addr, page_addr, endness, memory, **kwargs):
        """
        Small utility function for behavior which is duplicated in load

        mutates result to generate a new memory object and replace the last entry in it, which is None. Then, it will
        insert the new memory object into self.content.
        """
        global_end_addr = addr + page_addr
        global_start_addr = result[-1][0]
        size = global_end_addr - global_start_addr
        new_ast = self._default_value(global_start_addr, size, name='%s_%x' % (memory.id, global_start_addr), key=(self.category, global_start_addr), memory=memory, **kwargs)
        new_item = SimMemoryObject(new_ast, global_start_addr, endness=endness, byte_width=memory.state.arch.byte_width if memory is not None else 8)
        for subaddr in range(global_start_addr - page_addr, addr):
            self.content[subaddr] = new_item
        result[-1] = (global_start_addr, new_item)

    def store(self, addr, data, size=None, endness=None, memory=None, cooperate=False, **kwargs):
        if not cooperate:
            data = self._force_store_cooperation(addr, data, size, endness, memory=memory, **kwargs)

        if size == len(self.content) and addr == 0:
            self.sinkhole = data
            self.content = [None]*len(self.content)
        else:
            for subaddr in range(addr, addr + size):
                self.content[subaddr] = data

    def merge(self, others: List['ListPage'], merge_conditions, common_ancestor=None, page_addr: int=None,
              memory=None):

        changed_offsets = set()
        for other in others:
            changed_offsets |= self.changed_bytes(other, page_addr)

        all_pages: List['ListPage'] = [self] + others
        if merge_conditions is None:
            merge_conditions = [None] * len(all_pages)

        merged_to = None
        merged_objects = set()
        merged_offsets = set()
        for b in sorted(changed_offsets):
            if merged_to is not None and not b >= merged_to:
                l.info("merged_to = %d ... already merged byte 0x%x", merged_to, b)
                continue
            l.debug("... on byte 0x%x", b)

            memory_objects = []
            unconstrained_in = []

            # first get a list of all memory objects at that location, and
            # all memories that don't have those bytes
            for sm, fv in zip(all_pages, merge_conditions):
                if sm._contains(b, page_addr):
                    l.info("... present in %s", fv)
                    memory_objects.append((sm.content[b], fv))
                else:
                    l.info("... not present in %s", fv)
                    unconstrained_in.append((sm, fv))

            mos = set(mo for mo, _ in memory_objects)
            mo_bases = set(mo.base for mo, _ in memory_objects)
            mo_lengths = set(mo.length for mo, _ in memory_objects)

            if not unconstrained_in and not (mos - merged_objects):
                continue

            # first, optimize the case where we are dealing with the same-sized memory objects
            if len(mo_bases) == 1 and len(mo_lengths) == 1 and not unconstrained_in:
                our_mo = self.content[b]
                to_merge = [(mo.object, fv) for mo, fv in memory_objects]

                # Update `merged_to`
                mo_base = list(mo_bases)[0]
                mo_length = memory_objects[0][0].length
                size = mo_length - (page_addr + b - mo_base)
                merged_to = b + size

                merged_val = self._merge_values(to_merge, mo_length, memory=memory)
                if merged_val is None:
                    # merge_values() determines that we should not attempt to merge this value
                    continue

                # do the replacement
                # TODO: Implement in-place replacement instead of calling store()
                # new_object = self._replace_memory_object(our_mo, merged_val, page_addr, memory.page_size)

                self.store(b,
                           SimMemoryObject(merged_val, mo_base, endness='Iend_BE'),
                           size=size,
                           cooperate=True
                           )
                # merged_objects.add(new_object)
                # merged_objects.update(mos)
                merged_offsets.add(b)

            else:
                # get the size that we can merge easily. This is the minimum of
                # the size of all memory objects and unallocated spaces.
                min_size = min([mo.length - (b + page_addr - mo.base) for mo, _ in memory_objects])
                for um, _ in unconstrained_in:
                    for i in range(0, min_size):
                        if um._contains(b + i, page_addr):
                            min_size = i
                            break
                merged_to = b + min_size
                l.info("... determined minimum size of %d", min_size)

                # Now, we have the minimum size. We'll extract/create expressions of that
                # size and merge them
                extracted = [(mo.bytes_at(page_addr+b, min_size), fv) for mo, fv in memory_objects] if min_size != 0 else []
                created = [
                    (self._default_value(None, min_size, name="merge_uc_%s_%x" % (uc.id, b), memory=memory),
                     fv) for
                    uc, fv in unconstrained_in
                ]
                to_merge = extracted + created

                merged_val = self._merge_values(to_merge, min_size, memory=memory)
                if merged_val is None:
                    continue

                self.store(b,
                           SimMemoryObject(merged_val, page_addr+b, endness='Iend_BE'),
                           size=min_size,
                           endness='Iend_BE', cooperate=True
                           )  # do not convert endianness again
                merged_offsets.add(b)

        return merged_offsets

    def changed_bytes(self, other: 'ListPage', page_addr: int=None):

        candidates: Set[int] = set()
        for i in range(len(self.content)):
            if self._contains(i, page_addr):
                candidates.add(i)
        for i in range(len(other.content)):
            if other._contains(i, page_addr):
                candidates.add(i)

        byte_width = 8  # TODO: Introduce self.state if we want to use self.state.arch.byte_width
        differences: Set[int] = set()
        for c in candidates:
            if not self._contains(c, page_addr) and other._contains(c, page_addr):
                differences.add(c)
            elif self._contains(c, page_addr) and not other._contains(c, page_addr):
                differences.add(c)
            else:
                if self.content[c] is None:
                    self.content[c] = SimMemoryObject(self.sinkhole.bytes_at(page_addr+c, 1), page_addr + c,
                                                      byte_width=byte_width, endness='Iend_BE')
                if other.content[c] is None:
                    other.content[c] = SimMemoryObject(other.sinkhole.bytes_at(page_addr+c, 1), page_addr + c,
                                                       byte_width=byte_width, endness='Iend_BE')
                if self._contains(c, page_addr) and self.content[c] != other.content[c]:
                    # Try to see if the bytes are equal
                    self_byte = self.content[c].bytes_at(page_addr + c, 1)
                    other_byte = other.content[c].bytes_at(page_addr + c, 1)
                    if self_byte is not other_byte:
                        #l.debug("%s: offset %x, two different bytes %s %s from %s %s", self.id, c,
                        #        self_byte, other_byte,
                        #        self[c].object.model, other[c].object.model)
                        differences.add(c)
                else:
                    # this means the byte is in neither memory
                    pass

        return differences

    def _contains(self, off: int, page_addr: int):
        if off >= len(self.content):
            return False
        if self.content[off] is not None:
            return True
        if self.sinkhole is None:
            return False
        return self.sinkhole.includes(page_addr + off)

    def _replace_mo(self, old_mo: SimMemoryObject, new_mo: SimMemoryObject, page_addr: int,
                    page_size: int) -> SimMemoryObject:
        if self.sinkhole is old_mo:
            self.sinkhole = new_mo
        else:
            start, end = self._resolve_range(old_mo, page_addr, page_size)
            for i in range(start, end):
                if self.content[i-page_addr] is old_mo:
                    self.content[i-page_addr] = new_mo
        return new_mo

    @staticmethod
    def _resolve_range(mo: SimMemoryObject, page_addr: int, page_size) -> Tuple[int,int]:
        start = max(mo.base, page_addr)
        end = min(mo.last_addr + 1, page_addr + page_size)
        if end <= start:
            l.warning("Nothing left of the memory object to store in SimPage.")
        return start, end

    def _get_object(self, start: int, page_addr: int) -> Optional[SimMemoryObject]:
        mo = self.content[start]
        if mo is None:
            return None
        if mo.includes(start + page_addr):
            return mo
        return None
<|endoftext|>


<commit_before>import webapp2
import jinja2
import os
import urllib
import json
from google.appengine.api import users

from google.appengine.ext import ndb

html_dir = os.path.join(os.path.dirname(__file__), 'html')
jinja_env = jinja2.Environment(loader = jinja2.FileSystemLoader(html_dir), autoescape = True)

class Weather():
    def __init__(self, city, country, temperature, weatherType):
        self.city = city
        self.country = country
        self.temperature = temperature
        self.weatherType = weatherType

class Course():
    def __init__(self, name, links):
      self.name = name
      self.links = links

    def key(self):
      return ndb.Key('CourseTable', self.name).id();

class Handler(webapp2.RequestHandler):
    def write(self, *a, **kw):
        self.response.out.write(*a, **kw)

    def render_str(self, template, **params):
        t = jinja_env.get_template(template)
        return t.render(params)

    def render(self, template, **kw):
        self.write(self.render_str(template, **kw))

class TextbookTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    name = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class CourseTable(ndb.Model):
    name = ndb.StringProperty(required = True)
    courseId = ndb.StringProperty(required = True)

class CourseLinkTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    courseId = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class HousingTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class ProcrastinationTable(ndb.Model):
    link = ndb.StringProperty(required = True)

class MainPage(Handler):
    def render_with_data(self):
        data = urllib.urlopen("http://api.openweathermap.org/data/2.5/weather?q=Waterloo,ca")
        json_weather = json.loads(str(data.read()))
        city = json_weather["name"]
        country = json_weather["sys"]["country"]
        temperature = json_weather["main"]["temp"] - 273.15
        weatherType = json_weather["weather"][0]["main"]
        self.render("index.html", weather = Weather(city, country, temperature, weatherType))

    def get(self):
        self.render_with_data()

class Link():
    def __init__(self, link, courseId = None):
        self.link = link
        self.courseId = courseId

    def key(self, table):
        return ndb.Key(table, self.link).id()

    def Key(self, table):
        return ndb.Key(table, self.link)

class CoursesPage(Handler):
    def render_courses(self):
        course_query = CourseTable.query()
        courses_list = course_query.fetch()
        courses = []
        for course in courses_list:
            link_query = CourseLinkTable.query(CourseLinkTable.courseId==Course(course.name,[]).key())
            link_list = link_query.fetch()
            links = [Link(link.link) for link in link_list]
            courses.append(Course(course.name, links))
        self.render("courses.html", courses=courses)

    def get(self):
        self.render_courses()

    def post(self):
        course_val = self.request.get("addcourse")
        if(course_val):
            course = CourseTable(name = course_val, courseId = course_val)
            course.put()
            course.courseId = Course(course_val, []).key()
            course.put()
            self.render_courses()
        elif "addcourselink" in self.request.arguments()[0]:
            idd = self.request.arguments()[0]
            link_val = self.request.get(idd)
            if(link_val):
                idd = idd.replace("addcourselink", "")
                link = CourseLinkTable(link = link_val, courseId = idd, linkKey = link_val)
                link.put()
                link.linkKey = Link(link_val).key("CourseLinkTable")
                link.put()
                self.render_courses()
        elif "deletecourselink" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("deletecourselink", "")
            links = CourseLinkTable.query(CourseLinkTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_courses()
        elif "deletecourse" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("deletecourse", "")
            courses = CourseTable.query(CourseTable.courseId == remove_id)
            for course in courses:
                link_query = CourseLinkTable.query(CourseLinkTable.courseId==Course(course.name,[]).key())
                link_list = link_query.fetch()
                for link in link_list:
                    link.key.delete()
                course.key.delete()
            self.render_courses()

class TextbooksPage(Handler):
    def render_links(self):
        textbook_query = TextbookTable.query()
        link_list = textbook_query.fetch()
        links = [Link(link.link) for link in link_list]
        self.render("textbooks.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        isbn_num = self.request.get("addtextbook")
        link_val = "www.amazon.com/s/ref=nb_sb_noss?url=search-alias%3Daps&field-keywords=" + isbn_num
        data = urllib.urlopen("https://www.googleapis.com/books/v1/volumes?q=isbn:" + isbn_num)
        json_book = json.loads(str(data.read()))
        book_name = json_book["items"][0]["volumeInfo"]["title"]
        if(link_val):
            link = TextbookTable(link = link_val, name = book_name, linkKey = link_val)
            link.put()
            link.linkKey = Link(link_val).key("TextbookTable")
            link.put()
            self.render_links()
        elif "delete" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("delete", "")
            links = TextbookTable.query(TextbookTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_links()

class HousingPage(Handler):
    def render_links(self):
        housing_query = HousingTable.query()
        link_list = housing_query.fetch()
        links = [Link(link.link) for link in link_list]
        self.render("housing.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        link_val = self.request.get("addhousing")
        if(link_val):
            link = HousingTable(link = link_val, linkKey = link_val)
            link.put()
            link.linkKey = Link(link_val).key("HousingTable")
            link.put()
            self.render_links()
        elif "delete" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("delete", "")
            links = HousingTable.query(HousingTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_links()

class ProcrastinationPage(Handler):
    def render_links(self, latest=None):
        pro_query = ProcrastinationTable.query()
        link_list = pro_query.fetch()
        links = [link.link for link in link_list]
        if(latest):
            links.append(Link(latest.link))
        self.render("procrastination.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        link_val = self.request.get("addprocastination")
        if(link_val):
            link = ProcrastinationTable(link = link_val)
            link.put()
            self.render_links(link)

application = webapp2.WSGIApplication([
                  ('/', MainPage),
                  ('/courses', CoursesPage),
                  ('/textbooks', TextbooksPage),
                  ('/housing', HousingPage),
                  ('/procrastination', ProcrastinationPage)],
                  debug = True);
<commit_msg>Fixing merge issues with textbooks<commit_after>import webapp2
import jinja2
import os
import urllib
import json
from google.appengine.api import users

from google.appengine.ext import ndb

html_dir = os.path.join(os.path.dirname(__file__), 'html')
jinja_env = jinja2.Environment(loader = jinja2.FileSystemLoader(html_dir), autoescape = True)

class Weather():
    def __init__(self, city, country, temperature, weatherType):
        self.city = city
        self.country = country
        self.temperature = temperature
        self.weatherType = weatherType

class Course():
    def __init__(self, name, links):
      self.name = name
      self.links = links

    def key(self):
      return ndb.Key('CourseTable', self.name).id();

class Handler(webapp2.RequestHandler):
    def write(self, *a, **kw):
        self.response.out.write(*a, **kw)

    def render_str(self, template, **params):
        t = jinja_env.get_template(template)
        return t.render(params)

    def render(self, template, **kw):
        self.write(self.render_str(template, **kw))

class TextbookTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    name = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class CourseTable(ndb.Model):
    name = ndb.StringProperty(required = True)
    courseId = ndb.StringProperty(required = True)

class CourseLinkTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    courseId = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class HousingTable(ndb.Model):
    link = ndb.StringProperty(required = True)
    linkKey = ndb.StringProperty(required = True)

class ProcrastinationTable(ndb.Model):
    link = ndb.StringProperty(required = True)

class MainPage(Handler):
    def render_with_data(self):
        data = urllib.urlopen("http://api.openweathermap.org/data/2.5/weather?q=Waterloo,ca")
        json_weather = json.loads(str(data.read()))
        city = json_weather["name"]
        country = json_weather["sys"]["country"]
        temperature = json_weather["main"]["temp"] - 273.15
        weatherType = json_weather["weather"][0]["main"]
        self.render("index.html", weather = Weather(city, country, temperature, weatherType))

    def get(self):
        self.render_with_data()

class Link():
    def __init__(self, link, name, courseId = None):
        self.link = link
        self.courseId = courseId
        self.name = name

    def key(self, table):
        return ndb.Key(table, self.link).id()

class CoursesPage(Handler):
    def render_courses(self):
        course_query = CourseTable.query()
        courses_list = course_query.fetch()
        courses = []
        for course in courses_list:
            link_query = CourseLinkTable.query(CourseLinkTable.courseId==Course(course.name,[]).key())
            link_list = link_query.fetch()
            links = [Link(link.link) for link in link_list]
            courses.append(Course(course.name, links))
        self.render("courses.html", courses=courses)

    def get(self):
        self.render_courses()

    def post(self):
        course_val = self.request.get("addcourse")
        if(course_val):
            course = CourseTable(name = course_val, courseId = course_val)
            course.put()
            course.courseId = Course(course_val, []).key()
            course.put()
            self.render_courses()
        elif "addcourselink" in self.request.arguments()[0]:
            idd = self.request.arguments()[0]
            link_val = self.request.get(idd)
            if(link_val):
                idd = idd.replace("addcourselink", "")
                link = CourseLinkTable(link = link_val, courseId = idd, linkKey = link_val)
                link.put()
                link.linkKey = Link(link_val).key("CourseLinkTable")
                link.put()
                self.render_courses()
        elif "deletecourselink" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("deletecourselink", "")
            links = CourseLinkTable.query(CourseLinkTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_courses()
        elif "deletecourse" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("deletecourse", "")
            courses = CourseTable.query(CourseTable.courseId == remove_id)
            for course in courses:
                link_query = CourseLinkTable.query(CourseLinkTable.courseId==Course(course.name,[]).key())
                link_list = link_query.fetch()
                for link in link_list:
                    link.key.delete()
                course.key.delete()
            self.render_courses()

class TextbooksPage(Handler):
    def render_links(self):
        textbook_query = TextbookTable.query()
        link_list = textbook_query.fetch()
        links = [Link(link.link, link.name) for link in link_list]
        self.render("textbooks.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        isbn_num = self.request.get("addtextbook")
        if(isbn_num):
            link_val = "www.amazon.com/s/ref=nb_sb_noss?url=search-alias%3Daps&field-keywords=" + isbn_num
            data = urllib.urlopen("https://www.googleapis.com/books/v1/volumes?q=isbn:" + isbn_num)
            json_book = json.loads(str(data.read()))
            book_name = json_book["items"][0]["volumeInfo"]["title"]
            link = TextbookTable(link = link_val, name = book_name, linkKey = link_val)
            link.put()
            link.linkKey = Link(link_val, "").key("TextbookTable")
            link.put()
            self.render_links()
        elif "delete" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("delete", "")
            print remove_id
            links = TextbookTable.query(TextbookTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_links()

class HousingPage(Handler):
    def render_links(self):
        housing_query = HousingTable.query()
        link_list = housing_query.fetch()
        links = [Link(link.link) for link in link_list]
        self.render("housing.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        link_val = self.request.get("addhousing")
        if(link_val):
            link = HousingTable(link = link_val, linkKey = link_val)
            link.put()
            link.linkKey = Link(link_val).key("HousingTable")
            link.put()
            self.render_links()
        elif "delete" in self.request.arguments()[0]:
            remove_id = self.request.arguments()[0].replace("delete", "")
            links = HousingTable.query(HousingTable.linkKey == remove_id)
            for link in links:
       	        link.key.delete()
            self.render_links()

class ProcrastinationPage(Handler):
    def render_links(self, latest=None):
        pro_query = ProcrastinationTable.query()
        link_list = pro_query.fetch()
        links = [link.link for link in link_list]
        if(latest):
            links.append(Link(latest.link))
        self.render("procrastination.html", links=links)

    def get(self):
        self.render_links()

    def post(self):
        link_val = self.request.get("addprocastination")
        if(link_val):
            link = ProcrastinationTable(link = link_val)
            link.put()
            self.render_links(link)

application = webapp2.WSGIApplication([
                  ('/', MainPage),
                  ('/courses', CoursesPage),
                  ('/textbooks', TextbooksPage),
                  ('/housing', HousingPage),
                  ('/procrastination', ProcrastinationPage)],
                  debug = True);
<|endoftext|>


<commit_before>from __future__ import division, absolute_import
import os
import numpy as np
from scipy import spatial
import minipnm

''' Subclassing rules

    __init__ must *somehow* append coordinate (x,y,z) and connectivity (heads & tails)
    arrays to the network. that's its role. you can add many other attributes if
    you so wish, but you MUST provide those two.

    The other main rule is that aside from methods and properties, these objects
    should essentially be *just* dictionaries. One way to think about this is that
    the answer to the question "will I miss any data if I just write down the
    results of network.items?" should be NO.

    Convenience methods that automate some work (ie. import an image from a file)
    can be found in the `misc` module
'''

class Network(dict):
    filename = None

    @property
    def coords(self):
        return np.vstack([self['x'], self['y'], self['z']])

    @property
    def points(self):
        ''' a graph consists of edges and vertices
            in this *network* the points are vertices expressed in terms of the
            coordinates associated with them
        '''
        return np.vstack([self['x'], self['y'], self['z']]).T

    @property
    def pairs(self):
        ''' returns every 2-dimensional array representing an
            edge or connection between two points
        '''
        return np.vstack([self['heads'], self['tails']]).T

    @property
    def size(self):
        return len(self.points), len(self.pairs)

    @property
    def dims(self):
        w = self['x'].max() - self['x'].min()
        h = self['y'].max() - self['y'].min()
        t = self['z'].max() - self['z'].min() or 1
        return np.array([w, h, t])

    @property
    def centroid(self):
        x,y,z = self.coords
        w,h,t = self.dims
        return np.array([w/2.+x.min(), h/2.+y.min(), t/2.+z.min()])

    def boundary(self):
        all_points = np.arange(self.size[0])
        boundary_points = spatial.ConvexHull(self.points).vertices
        return np.in1d(all_points, boundary_points).astype(bool)

    def save(self, filename):
        minipnm.save_vtp(self, filename)
        self.filename = filename

    def render(self, values=[]):
        minipnm.render(self, values)

    def merge(self, other, axis=2, spacing=0, centering=True, stitch=True):
        new = Network()

        # alignment along a common centroid
        if centering:
            center_distance = other.centroid - self.centroid
            center_distance[axis] = 0 # we take care of this one differently
            shifted_points = other.points - center_distance
        else:
            shifted_points = other.points

        # the distance between the max for base and min for other should
        # equal to spacing. rearranging, it gives us the required offset
        offset = other.coords[axis].min() \
                - self.coords[axis].max() - spacing
        shifted_points.T[axis] -= offset
        new['x'], new['y'], new['z'] = np.vstack([self.points, shifted_points]).T

        # push the connectivity array by the number of already existing vertices
        Va, Ea = self.size
        new['heads'], new['tails'] = np.vstack([self.pairs,
                                                other.pairs+Va
                                                ]).T

        # merge the rest
        for key in set(self.keys()+other.keys()) - {'x','y','z','heads','tails'}:
            values_self = self.get(key, np.zeros(self.size[0]))
            values_other = other.get(key, np.zeros(other.size[0]))
            new[key] = np.hstack([values_self, values_other])

        return new

    def cut(self, mask):
        ''' returns id of throats where the head and the tail are on opposite
            sides of the mask
        '''
        return np.array([i for i, (hix, tix) in enumerate(
            zip(self['heads'], self['tails']))
            if ( mask[hix] and ~mask[tix])
            or (~mask[hix] and  mask[tix])
            ])

    def prune(self, inaccessible, remove_pores=True):
        new = self.copy()

        accessible = np.arange(new.size[0])[~inaccessible]
        good_heads = np.in1d(self['heads'], accessible)
        good_tails = np.in1d(self['tails'], accessible)
        new['heads'] = self['heads'][good_heads & good_tails]
        new['tails'] = self['tails'][good_heads & good_tails]
        if not remove_pores:
            return new

        # now we need to update any other values to the new indeces
        translate = dict(zip(accessible, np.arange(accessible.size)))
        new['x'] = new['x'][accessible]
        new['y'] = new['y'][accessible]
        new['z'] = new['z'][accessible]
        new['heads'] = np.array(map(translate.get, new['heads']))
        new['tails'] = np.array(map(translate.get, new['tails']))
        for key, array in self.items():
            if array.size == self.size[0]:
                new[key] = array[accessible]

        return new

    def copy(self):
        clone = self.__class__.__new__(self.__class__)
        clone.update(self)
        return clone

    def __repr__(self):
        return self.__class__.__name__+str(self.size)

    def __str__(self):
        entries = [self.__class__.__name__]
        for key, value in sorted(self.items()):
            entries.append('{:<15}: {:<10}: {:<15}'.format(
                    key,
                    value.dtype,
                    value.shape,))
        return '\n\t'.join(entries)

    def __add__(self, other):
        return self.merge(other)

    def __sub__(self, inaccessible):
        return self.prune(inaccessible)

class Cubic(Network):

    def __init__(self, ndarray, dims=[1,1,1]):
        self['intensity'] = ndarray.ravel()

        rel_coords = np.array(
            [idx for idx,val in np.ndenumerate(ndarray)]).astype(float)
        abs_coords = rel_coords * dims / np.array(np.subtract(ndarray.shape,1)).clip(1, np.inf)
        self['x'], self['y'], self['z'] = abs_coords.T

        I = np.arange(ndarray.size).reshape(ndarray.shape)
        heads, tails = [], []
        for A,B in [
            (I[:,:,:-1], I[:,:,1:]),
            (I[:,:-1], I[:,1:]),
            (I[:-1], I[1:]),
            ]:
            hs, ts = np.vstack([A.flat, B.flat])
            heads.extend(hs)
            tails.extend(ts)

        self['heads'] = np.array(heads)
        self['tails'] = np.array(tails)

    @property
    def resolution(self):
        return np.array([len(set(d)) for d in self.coords])

    def asarray(self, values=None):
        _ndarray = np.zeros(self.resolution)
        rel_coords = np.true_divide(self.points, self.dims)*(self.resolution-1)
        rel_coords = np.rint(rel_coords).astype(int) # absolutely bizarre bug

        actual_indexes = np.ravel_multi_index(rel_coords.T, self.resolution)
        if values==None:
            values = self['intensity']
        _ndarray.flat[actual_indexes] = values
        return _ndarray

class Delaunay(Network):

    def __init__(self, points, mask=None):
        self['x'], self['y'], self['z'] = points.T
        self['heads'], self['tails'] = self.edges_from_points(points).T

    @staticmethod
    def edges_from_points(points, mask=None):
        triangulation = spatial.Delaunay(points)
        edges = set()
        for a,b,c,d in triangulation.vertices:
            if mask is not None:
                a,b,c,d = (mask[i] for i in [a,b,c,d])
            for edge in [(a,b),(b,c),(c,d),(d,a)]:
                edge = tuple(sorted(edge))
                if edge not in edges:
                    edges.add(edge)

        return np.array(list(edges))<commit_msg>Added 'flux' method to network<commit_after>from __future__ import division, absolute_import
import os
import numpy as np
from scipy import spatial
import minipnm

''' Subclassing rules

    __init__ must *somehow* append coordinate (x,y,z) and connectivity (heads & tails)
    arrays to the network. that's its role. you can add many other attributes if
    you so wish, but you MUST provide those two.

    The other main rule is that aside from methods and properties, these objects
    should essentially be *just* dictionaries. One way to think about this is that
    the answer to the question "will I miss any data if I just write down the
    results of network.items?" should be NO.

    Convenience methods that automate some work (ie. import an image from a file)
    can be found in the `misc` module
'''

class Network(dict):
    filename = None

    @property
    def coords(self):
        return np.vstack([self['x'], self['y'], self['z']])

    @property
    def points(self):
        ''' a graph consists of edges and vertices
            in this *network* the points are vertices expressed in terms of the
            coordinates associated with them
        '''
        return np.vstack([self['x'], self['y'], self['z']]).T

    @property
    def pairs(self):
        ''' returns every 2-dimensional array representing an
            edge or connection between two points
        '''
        return np.vstack([self['heads'], self['tails']]).T

    @property
    def size(self):
        return len(self.points), len(self.pairs)

    @property
    def dims(self):
        w = self['x'].max() - self['x'].min()
        h = self['y'].max() - self['y'].min()
        t = self['z'].max() - self['z'].min() or 1
        return np.array([w, h, t])

    @property
    def centroid(self):
        x,y,z = self.coords
        w,h,t = self.dims
        return np.array([w/2.+x.min(), h/2.+y.min(), t/2.+z.min()])

    def boundary(self):
        all_points = np.arange(self.size[0])
        boundary_points = spatial.ConvexHull(self.points).vertices
        return np.in1d(all_points, boundary_points).astype(bool)

    def save(self, filename):
        minipnm.save_vtp(self, filename)
        self.filename = filename

    def render(self, values=[]):
        minipnm.render(self, values)

    def merge(self, other, axis=2, spacing=0, centering=True, stitch=True):
        new = Network()

        # alignment along a common centroid
        if centering:
            center_distance = other.centroid - self.centroid
            center_distance[axis] = 0 # we take care of this one differently
            shifted_points = other.points - center_distance
        else:
            shifted_points = other.points

        # the distance between the max for base and min for other should
        # equal to spacing. rearranging, it gives us the required offset
        offset = other.coords[axis].min() \
                - self.coords[axis].max() - spacing
        shifted_points.T[axis] -= offset
        new['x'], new['y'], new['z'] = np.vstack([self.points, shifted_points]).T

        # push the connectivity array by the number of already existing vertices
        Va, Ea = self.size
        new['heads'], new['tails'] = np.vstack([self.pairs,
                                                other.pairs+Va
                                                ]).T

        # merge the rest
        for key in set(self.keys()+other.keys()) - {'x','y','z','heads','tails'}:
            values_self = self.get(key, np.zeros(self.size[0]))
            values_other = other.get(key, np.zeros(other.size[0]))
            new[key] = np.hstack([values_self, values_other])

        return new

    def cut(self, mask):
        ''' returns id of throats where the head and the tail are on opposite
            sides of the mask
        '''
        return np.array([i for i, (hix, tix) in enumerate(
            zip(self['heads'], self['tails']))
            if ( mask[hix] and ~mask[tix])
            or (~mask[hix] and  mask[tix])
            ])

    def flux(self, data, cut_ids):
        vertex_ids = self.pairs[cut_ids].flatten()
        data_pairs = data[vertex_ids].reshape(vertex_ids.size//2, 2)
        data_delta = np.abs(np.diff(data_pairs))
        return data_delta

    def prune(self, inaccessible, remove_pores=True):
        new = self.copy()

        accessible = np.arange(new.size[0])[~inaccessible]
        good_heads = np.in1d(self['heads'], accessible)
        good_tails = np.in1d(self['tails'], accessible)
        new['heads'] = self['heads'][good_heads & good_tails]
        new['tails'] = self['tails'][good_heads & good_tails]
        if not remove_pores:
            return new

        # now we need to update any other values to the new indeces
        translate = dict(zip(accessible, np.arange(accessible.size)))
        new['x'] = new['x'][accessible]
        new['y'] = new['y'][accessible]
        new['z'] = new['z'][accessible]
        new['heads'] = np.array(map(translate.get, new['heads']))
        new['tails'] = np.array(map(translate.get, new['tails']))
        for key, array in self.items():
            if array.size == self.size[0]:
                new[key] = array[accessible]

        return new

    def copy(self):
        clone = self.__class__.__new__(self.__class__)
        clone.update(self)
        return clone

    def __repr__(self):
        return self.__class__.__name__+str(self.size)

    def __str__(self):
        entries = [self.__class__.__name__]
        for key, value in sorted(self.items()):
            entries.append('{:<15}: {:<10}: {:<15}'.format(
                    key,
                    value.dtype,
                    value.shape,))
        return '\n\t'.join(entries)

    def __add__(self, other):
        return self.merge(other)

    def __sub__(self, inaccessible):
        return self.prune(inaccessible)

class Cubic(Network):

    def __init__(self, ndarray, dims=[1,1,1]):
        self['intensity'] = ndarray.ravel()

        rel_coords = np.array(
            [idx for idx,val in np.ndenumerate(ndarray)]).astype(float)
        abs_coords = rel_coords * dims / np.array(np.subtract(ndarray.shape,1)).clip(1, np.inf)
        self['x'], self['y'], self['z'] = abs_coords.T

        I = np.arange(ndarray.size).reshape(ndarray.shape)
        heads, tails = [], []
        for A,B in [
            (I[:,:,:-1], I[:,:,1:]),
            (I[:,:-1], I[:,1:]),
            (I[:-1], I[1:]),
            ]:
            hs, ts = np.vstack([A.flat, B.flat])
            heads.extend(hs)
            tails.extend(ts)

        self['heads'] = np.array(heads)
        self['tails'] = np.array(tails)

    @property
    def resolution(self):
        return np.array([len(set(d)) for d in self.coords])

    def asarray(self, values=None):
        _ndarray = np.zeros(self.resolution)
        rel_coords = np.true_divide(self.points, self.dims)*(self.resolution-1)
        rel_coords = np.rint(rel_coords).astype(int) # absolutely bizarre bug

        actual_indexes = np.ravel_multi_index(rel_coords.T, self.resolution)
        if values==None:
            values = self['intensity']
        _ndarray.flat[actual_indexes] = values
        return _ndarray

class Delaunay(Network):

    def __init__(self, points, mask=None):
        self['x'], self['y'], self['z'] = points.T
        self['heads'], self['tails'] = self.edges_from_points(points).T

    @staticmethod
    def edges_from_points(points, mask=None):
        triangulation = spatial.Delaunay(points)
        edges = set()
        for a,b,c,d in triangulation.vertices:
            if mask is not None:
                a,b,c,d = (mask[i] for i in [a,b,c,d])
            for edge in [(a,b),(b,c),(c,d),(d,a)]:
                edge = tuple(sorted(edge))
                if edge not in edges:
                    edges.add(edge)

        return np.array(list(edges))

if __name__ == '__main__':
    R = np.ones([30,30,1])

    network = Cubic(R)
    x,y,z = network.coords
    cut = network.cut(x<0.5)
    ids = network.pairs[cut].flatten()
    print x[ids].reshape(ids.size//2, 2)<|endoftext|>


<commit_before>from hail.utils import warn, error, java
from .indices import *
from ..expressions import Expression, Aggregable, ExpressionException, expr_any
from typing import *

@typecheck(caller=str,
           expr=Expression,
           expected_indices=Indices,
           aggregation_axes=setof(str),
           broadcast=bool)
def analyze(caller: str,
            expr: Expression,
            expected_indices: Indices,
            aggregation_axes: Set = set(),
            broadcast=True):
    indices = expr._indices
    source = indices.source
    axes = indices.axes
    aggregations = expr._aggregations

    warnings = []
    errors = []

    expected_source = expected_indices.source
    expected_axes = expected_indices.axes

    if source is not None and source is not expected_source:
        bad_refs = []
        for name, inds in get_refs(expr).items():
            if inds.source is not expected_source:
                bad_refs.append(name)
        errors.append(
            ExpressionException("'{caller}': source mismatch\n"
                                "  Expected an expression from source {expected}\n"
                                "  Found expression derived from source {actual}\n"
                                "  Problematic field(s): {bad_refs}\n\n"
                                "  This error is commonly caused by chaining methods together:\n"
                                "    >>> ht.distinct().select(ht.x)\n\n"
                                "  Correct usage:\n"
                                "    >>> ht = ht.distinct()\n"
                                "    >>> ht = ht.select(ht.x)".format(
                caller=caller,
                expected=expected_source,
                actual=source,
                bad_refs=list(bad_refs)
            )))

    # check for stray indices by subtracting expected axes from observed
    if broadcast:
        unexpected_axes = axes - expected_axes
        strictness = ''
    else:
        unexpected_axes = axes if axes != expected_axes else set()
        strictness = 'strictly '

    if unexpected_axes:
        # one or more out-of-scope fields
        refs = get_refs(expr)
        bad_refs = []
        for name, inds in refs.items():
            if broadcast:
                bad_axes = inds.axes.intersection(unexpected_axes)
                if bad_axes:
                    bad_refs.append((name, inds))
            else:
                if inds.axes != expected_axes:
                    bad_refs.append((name, inds))

        assert len(bad_refs) > 0
        errors.append(ExpressionException(
            "scope violation: '{caller}' expects an expression {strictness}indexed by {expected}"
            "\n    Found indices {axes}, with unexpected indices {stray}. Invalid fields:{fields}{agg}".format(
                caller=caller,
                strictness=strictness,
                expected=list(expected_axes),
                axes=list(indices.axes),
                stray=list(unexpected_axes),
                fields=''.join("\n        '{}' (indices {})".format(name, list(inds.axes)) for name, inds in bad_refs),
                agg='' if (unexpected_axes - aggregation_axes) else
                "\n    '{}' supports aggregation over axes {}, "
                "so these fields may appear inside an aggregator function.".format(caller, list(aggregation_axes))
            )))

    if aggregations:
        if aggregation_axes:

            # the expected axes of aggregated expressions are the expected axes + axes aggregated over
            expected_agg_axes = expected_axes.union(aggregation_axes)

            for agg in aggregations:
                assert isinstance(agg, Aggregation)
                refs = get_refs(*agg.exprs)
                agg_indices = agg.indices
                agg_axes = agg_indices.axes

                # check for stray indices
                unexpected_agg_axes = agg_axes - expected_agg_axes
                if unexpected_agg_axes:
                    # one or more out-of-scope fields
                    bad_refs = []
                    for name, inds in refs.items():
                        bad_axes = inds.axes.intersection(unexpected_agg_axes)
                        if bad_axes:
                            bad_refs.append((name, inds))

                    assert len(bad_refs) > 0

                    errors.append(ExpressionException(
                        "scope violation: '{caller}' supports aggregation over indices {expected}"
                        "\n    Found indices {axes}, with unexpected indices {stray}. Invalid fields:{fields}".format(
                            caller=caller,
                            expected=list(aggregation_axes),
                            axes=list(agg_axes),
                            stray=list(unexpected_agg_axes),
                            fields=''.join("\n        '{}' (indices {})".format(
                                name, list(inds.axes)) for name, inds in bad_refs)
                        )
                    ))
        else:
            errors.append(ExpressionException("'{}' does not support aggregation".format(caller)))

    for w in warnings:
        warn('{}'.format(w.msg))
    if errors:
        for e in errors:
            error('{}'.format(e.msg))
        raise errors[0]

@typecheck(expression=expr_any)
def eval_expr(expression):
    """Evaluate a Hail expression, returning the result.

    This method is extremely useful for learning about Hail expressions and understanding
    how to compose them.

    The expression must have no indices, but can refer to the globals
    of a :class:`.hail.Table` or :class:`.hail.MatrixTable`.

    Examples
    --------

    Evaluate a conditional:

    >>> x = 6
    >>> hl.eval_expr(hl.cond(x % 2 == 0, 'Even', 'Odd'))
    'Even'

    Parameters
    ----------
    expression : :class:`.Expression`
        Any expression, or a Python value that can be implicitly interpreted as an expression.

    Returns
    -------
    any
        Result of evaluating `expression`.
    """
    return eval_expr_typed(expression)[0]

@typecheck(expression=expr_any)
def eval_expr_typed(expression):
    """Evaluate a Hail expression, returning the result and the type of the result.

    This method is extremely useful for learning about Hail expressions and understanding
    how to compose them.

    The expression must have no indices, but can refer to the globals
    of a :class:`.hail.Table` or :class:`.hail.MatrixTable`.

    Examples
    --------
    Evaluate a conditional:

    >>> x = 6
    >>> hl.eval_expr_typed(hl.cond(x % 2 == 0, 'Even', 'Odd'))
    ('Odd', tstr)

    Parameters
    ----------
    expression : :class:`.Expression`
        Any expression, or a Python value that can be implicitly interpreted as an expression.

    Returns
    -------
    (any, :class:`.HailType`)
        Result of evaluating `expression`, and its type.

    """
    analyze('eval_expr_typed', expression, Indices(expression._indices.source))

    if expression._indices.source is None:
        return (expression.dtype._from_json(
            java.Env._hail_package.expr.ir.Interpret.interpretPyIR(str(expression._ir))),
                expression.dtype)
    else:
        return expression.collect()[0], expression.dtype

def _get_refs(expr: Union[Expression, Aggregable], builder: Dict[str, Indices]) -> None:
    from hail.ir import GetField, TopLevelReference

    for ir in expr._ir.search(
            lambda a: isinstance(a, GetField)
                      and not a.name.startswith('__uid')
                      and isinstance(a.o, TopLevelReference)):
        src = expr._indices.source
        builder[ir.name] = src._indices_from_ref[ir.o.name]

def extract_refs_by_indices(exprs, indices):
    """Returns a set of references in `exprs` with indices `indices`.

    Parameters
    ----------
    expr : Expression
    indices : Indices

    Returns
    -------
    Set[str]
    """
    s = set()
    for e in exprs:
        for name, inds in get_refs(e).items():
            if inds == indices:
                s.add(name)
    return s

def get_refs(*exprs: Union[Expression, Aggregable]) -> Dict[str, Indices]:
    builder = {}
    for e in exprs:
        _get_refs(e, builder)
    return builder

@typecheck(caller=str,
           expr=Expression)
def matrix_table_source(caller, expr):
    from hail import MatrixTable
    source = expr._indices.source
    if not isinstance(source, MatrixTable):
        raise ValueError(
            "{}: Expect an expression of 'MatrixTable', found {}".format(
                caller,
                "expression of '{}'".format(source.__class__) if source is not None else 'scalar expression'))
    return source

@typecheck(caller=str,
           expr=Expression)
def check_entry_indexed(caller, expr):
    if expr._indices.source is None:
        raise ExpressionException("{}: expression must be entry-indexed,"
                                  " found no indices (no source)")
    if expr._indices != expr._indices.source._entry_indices:
        raise ExpressionException("{}: expression must be entry-indexed,"
                                  " found indices {}".format(caller, list(expr._indices.axes)))

@typecheck(caller=str,
           expr=Expression)
def check_row_indexed(caller, expr):
    if expr._indices.source is None:
        raise ExpressionException("{}: expression must be row-indexed,"
                                  " found no indices (no source)")
    if expr._indices != expr._indices.source._row_indices:
        raise ExpressionException("{}: expression must be row-indexed,"
                                  " found indices {}".format(caller, list(expr._indices.axes)))
<commit_msg>Fix usage of _hail_package in eval_expr (#4182)<commit_after>from hail.utils import warn, error, java
from .indices import *
from ..expressions import Expression, Aggregable, ExpressionException, expr_any
from typing import *

@typecheck(caller=str,
           expr=Expression,
           expected_indices=Indices,
           aggregation_axes=setof(str),
           broadcast=bool)
def analyze(caller: str,
            expr: Expression,
            expected_indices: Indices,
            aggregation_axes: Set = set(),
            broadcast=True):
    indices = expr._indices
    source = indices.source
    axes = indices.axes
    aggregations = expr._aggregations

    warnings = []
    errors = []

    expected_source = expected_indices.source
    expected_axes = expected_indices.axes

    if source is not None and source is not expected_source:
        bad_refs = []
        for name, inds in get_refs(expr).items():
            if inds.source is not expected_source:
                bad_refs.append(name)
        errors.append(
            ExpressionException("'{caller}': source mismatch\n"
                                "  Expected an expression from source {expected}\n"
                                "  Found expression derived from source {actual}\n"
                                "  Problematic field(s): {bad_refs}\n\n"
                                "  This error is commonly caused by chaining methods together:\n"
                                "    >>> ht.distinct().select(ht.x)\n\n"
                                "  Correct usage:\n"
                                "    >>> ht = ht.distinct()\n"
                                "    >>> ht = ht.select(ht.x)".format(
                caller=caller,
                expected=expected_source,
                actual=source,
                bad_refs=list(bad_refs)
            )))

    # check for stray indices by subtracting expected axes from observed
    if broadcast:
        unexpected_axes = axes - expected_axes
        strictness = ''
    else:
        unexpected_axes = axes if axes != expected_axes else set()
        strictness = 'strictly '

    if unexpected_axes:
        # one or more out-of-scope fields
        refs = get_refs(expr)
        bad_refs = []
        for name, inds in refs.items():
            if broadcast:
                bad_axes = inds.axes.intersection(unexpected_axes)
                if bad_axes:
                    bad_refs.append((name, inds))
            else:
                if inds.axes != expected_axes:
                    bad_refs.append((name, inds))

        assert len(bad_refs) > 0
        errors.append(ExpressionException(
            "scope violation: '{caller}' expects an expression {strictness}indexed by {expected}"
            "\n    Found indices {axes}, with unexpected indices {stray}. Invalid fields:{fields}{agg}".format(
                caller=caller,
                strictness=strictness,
                expected=list(expected_axes),
                axes=list(indices.axes),
                stray=list(unexpected_axes),
                fields=''.join("\n        '{}' (indices {})".format(name, list(inds.axes)) for name, inds in bad_refs),
                agg='' if (unexpected_axes - aggregation_axes) else
                "\n    '{}' supports aggregation over axes {}, "
                "so these fields may appear inside an aggregator function.".format(caller, list(aggregation_axes))
            )))

    if aggregations:
        if aggregation_axes:

            # the expected axes of aggregated expressions are the expected axes + axes aggregated over
            expected_agg_axes = expected_axes.union(aggregation_axes)

            for agg in aggregations:
                assert isinstance(agg, Aggregation)
                refs = get_refs(*agg.exprs)
                agg_indices = agg.indices
                agg_axes = agg_indices.axes

                # check for stray indices
                unexpected_agg_axes = agg_axes - expected_agg_axes
                if unexpected_agg_axes:
                    # one or more out-of-scope fields
                    bad_refs = []
                    for name, inds in refs.items():
                        bad_axes = inds.axes.intersection(unexpected_agg_axes)
                        if bad_axes:
                            bad_refs.append((name, inds))

                    assert len(bad_refs) > 0

                    errors.append(ExpressionException(
                        "scope violation: '{caller}' supports aggregation over indices {expected}"
                        "\n    Found indices {axes}, with unexpected indices {stray}. Invalid fields:{fields}".format(
                            caller=caller,
                            expected=list(aggregation_axes),
                            axes=list(agg_axes),
                            stray=list(unexpected_agg_axes),
                            fields=''.join("\n        '{}' (indices {})".format(
                                name, list(inds.axes)) for name, inds in bad_refs)
                        )
                    ))
        else:
            errors.append(ExpressionException("'{}' does not support aggregation".format(caller)))

    for w in warnings:
        warn('{}'.format(w.msg))
    if errors:
        for e in errors:
            error('{}'.format(e.msg))
        raise errors[0]

@typecheck(expression=expr_any)
def eval_expr(expression):
    """Evaluate a Hail expression, returning the result.

    This method is extremely useful for learning about Hail expressions and understanding
    how to compose them.

    The expression must have no indices, but can refer to the globals
    of a :class:`.hail.Table` or :class:`.hail.MatrixTable`.

    Examples
    --------

    Evaluate a conditional:

    >>> x = 6
    >>> hl.eval_expr(hl.cond(x % 2 == 0, 'Even', 'Odd'))
    'Even'

    Parameters
    ----------
    expression : :class:`.Expression`
        Any expression, or a Python value that can be implicitly interpreted as an expression.

    Returns
    -------
    any
        Result of evaluating `expression`.
    """
    return eval_expr_typed(expression)[0]

@typecheck(expression=expr_any)
def eval_expr_typed(expression):
    """Evaluate a Hail expression, returning the result and the type of the result.

    This method is extremely useful for learning about Hail expressions and understanding
    how to compose them.

    The expression must have no indices, but can refer to the globals
    of a :class:`.hail.Table` or :class:`.hail.MatrixTable`.

    Examples
    --------
    Evaluate a conditional:

    >>> x = 6
    >>> hl.eval_expr_typed(hl.cond(x % 2 == 0, 'Even', 'Odd'))
    ('Odd', tstr)

    Parameters
    ----------
    expression : :class:`.Expression`
        Any expression, or a Python value that can be implicitly interpreted as an expression.

    Returns
    -------
    (any, :class:`.HailType`)
        Result of evaluating `expression`, and its type.

    """
    analyze('eval_expr_typed', expression, Indices(expression._indices.source))

    if expression._indices.source is None:
        return (expression.dtype._from_json(
            java.Env.hail().expr.ir.Interpret.interpretPyIR(str(expression._ir))),
                expression.dtype)
    else:
        return expression.collect()[0], expression.dtype

def _get_refs(expr: Union[Expression, Aggregable], builder: Dict[str, Indices]) -> None:
    from hail.ir import GetField, TopLevelReference

    for ir in expr._ir.search(
            lambda a: isinstance(a, GetField)
                      and not a.name.startswith('__uid')
                      and isinstance(a.o, TopLevelReference)):
        src = expr._indices.source
        builder[ir.name] = src._indices_from_ref[ir.o.name]

def extract_refs_by_indices(exprs, indices):
    """Returns a set of references in `exprs` with indices `indices`.

    Parameters
    ----------
    expr : Expression
    indices : Indices

    Returns
    -------
    Set[str]
    """
    s = set()
    for e in exprs:
        for name, inds in get_refs(e).items():
            if inds == indices:
                s.add(name)
    return s

def get_refs(*exprs: Union[Expression, Aggregable]) -> Dict[str, Indices]:
    builder = {}
    for e in exprs:
        _get_refs(e, builder)
    return builder

@typecheck(caller=str,
           expr=Expression)
def matrix_table_source(caller, expr):
    from hail import MatrixTable
    source = expr._indices.source
    if not isinstance(source, MatrixTable):
        raise ValueError(
            "{}: Expect an expression of 'MatrixTable', found {}".format(
                caller,
                "expression of '{}'".format(source.__class__) if source is not None else 'scalar expression'))
    return source

@typecheck(caller=str,
           expr=Expression)
def check_entry_indexed(caller, expr):
    if expr._indices.source is None:
        raise ExpressionException("{}: expression must be entry-indexed,"
                                  " found no indices (no source)")
    if expr._indices != expr._indices.source._entry_indices:
        raise ExpressionException("{}: expression must be entry-indexed,"
                                  " found indices {}".format(caller, list(expr._indices.axes)))

@typecheck(caller=str,
           expr=Expression)
def check_row_indexed(caller, expr):
    if expr._indices.source is None:
        raise ExpressionException("{}: expression must be row-indexed,"
                                  " found no indices (no source)")
    if expr._indices != expr._indices.source._row_indices:
        raise ExpressionException("{}: expression must be row-indexed,"
                                  " found indices {}".format(caller, list(expr._indices.axes)))
<|endoftext|>