<commit_before># Copyright (C) 2014 Diane Trout
#
# This package is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License Version 2.1
# as published by the Free Software Foundation or any newer version.

# This package is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
# License Version 2.1 for more details.

# You should have received a copy of the GNU Lesser General Public
# License Version 2.1 along with this package; if not, write to the
# Free Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston,
# MA 02110-1301 USA

import os
import re
import six

from six.moves import range
from six.moves.urllib.error import HTTPError
from six.moves.urllib.parse import urlparse
from six.moves.urllib.request import urlopen

import RDF
import collections

from IPython.core.magic import (Magics, magics_class, cell_magic, line_magic,
                                line_cell_magic, UsageError)
from IPython.core.magic_arguments import (
    argument, magic_arguments, parse_argstring)

_prefixes = {
    'rdf': "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
    'rdfs': "http://www.w3.org/2000/01/rdf-schema#",
    'owl': 'http://www.w3.org/2002/07/owl#',
    'dc': "http://purl.org/dc/elements/1.1/",
    'xml': 'http://www.w3.org/XML/1998/namespace',
    'xsd': "http://www.w3.org/2001/XMLSchema#",
    'vs': 'http://www.w3.org/2003/06/sw-vocab-status/ns#',
    'wot': 'http://xmlns.com/wot/0.1/'
}
_namespaces = dict(( (v, k) for k, v in _prefixes.items()))

def dump_model(model):
    export = RDF.Serializer(name='turtle')
    print (export.serialize_model_to_string(model))

def display_node(node, mime_type):
    if mime_type == 'text/html':
        template = '{curie}'
    else:
        template = '{curie}'

    if isinstance(node, RDF.Node) and node.is_resource():
        match = ""
        uri = str(node.uri)
        curie = uri
        for namespace, prefix in _namespaces.items():
            if uri.startswith(namespace) and len(namespace) > len(match):
                match = namespace

        if len(match) > 0:
            prefix = _namespaces[match]
            curie = uri.replace(match, "{0}:".format(prefix))

        return template.format(anchor=uri, curie=curie)
    else:
        return str(node)

class Node(RDF.Node):
    def _repr_html_(self):
        return display_node(self, mime_type='text/html')

Uri = RDF.Uri
NS = RDF.NS

class LibRdfResults(collections.Sequence):
    results = None
    columns = None
    def __init__(self, result_set):
        if result_set is not None:
            self.columns = self.get_bindings(result_set)
            self.results = self.get_results(result_set)

    def get_bindings(self, result_set):
        b = []
        for i in range(result_set.get_bindings_count()):
            b.append(result_set.get_binding_name(i))
        return b

    def get_results(self, result_set):
        results = []
        for row in result_set:
            record = [ row[k] for k in self.columns ]
            results.append(record)
        return results

    def __getitem__(self, key):
        return dict(zip(self.columns, self.results[key]))

    def __len__(self):
        return len(self.results)

    def generate_html(self):
        result_set = iter(self)
        yield '

'
        for c in self.columns:
            yield "{0}	".format(str(c))
        yield '
'
        for row in self:
            yield '
'
            for key in self.columns:
                value = row[key]
                yield ''
                yield display_node(value, mime_type='text/html')
                yield '	'
            yield '
'
        yield '
'

    def _repr_html_(self):
        return "".join(self.generate_html())

    def __str__(self):
        result_set = iter(self)
        output = []
        output.append("\t".join(self.columns))
        for row in self.results:
            ordered = (display_node(value, mime_type='text/plain')
                       for value in row)
            output.append("\t".join(ordered))
        return os.linesep.join(output)

@magics_class
class SPARQLMagics(Magics):
    SUPPORT_FROM = False
    def __init__(self, shell):
        super(SPARQLMagics, self).__init__(shell)

    @magic_arguments()
    @argument('prefix', nargs=1, type=str,
              help="Specify namespace prefix")
    @argument('namespace', nargs=1, type=str,
              help="specify namespace")
    @line_magic
    def addns(self, line):
        args = parse_argstring(self.addns, line)
        if args.prefix is None or args.namespace is None:
            raise UsageError("Please specify prefix and namespace")
        namespace = args.namespace[0]
        prefix = args.prefix[0]
        _namespaces[namespace] = prefix
        _prefixes[prefix] = namespace
        self.shell.user_ns[prefix] = RDF.NS(namespace)

    @line_magic
    def lsns(self, line=None):
        keys = sorted(_prefixes.keys())
        maxlen = max((len(k) for k in keys))
        for k in keys:
            print ("{prefix:{width}} {ns}".format(
                prefix=k,
                width=maxlen,
                ns=_prefixes[k]))

    @magic_arguments()
    @argument('name', nargs=1, type=str,
              help="remove an item from our default namespaces")
    @line_magic
    def delns(self, line=None):
        args = parse_argstring(self.delns, line)
        if args.name is None:
            raise UsageError("Specify what to delete")

        name = args.name[0]
        if name in _prefixes:
            del _namespaces[self._prefixes[name]]
            del _prefixes[name]
            del self.shell.user_ns[name]
        elif name in _namespaces:
            prefix = self._namespaces[name]
            del _prefixes[prefix]
            del _namespaces[name]
            del self.shell.user_ns[prefix]
        else:
            raise UsageError("%s was not found in our namespace cache" % (
                name,))

    @magic_arguments()
    @argument('-D', '--data', default=None, type=str,
              help="Add RDF Data source URI")
    @argument('-p', '--protocol', type=str, default=None,
              help="Call a SPARQL HTTP Protocol at service URI")
    @argument('-m', '--model', type=str, default=None,
              help="Specify a model to use")
    @argument('--output', type=str,
              default=None,
              help="Specifiy variable to hold output")
    @line_cell_magic
    def roqet(self, line, cell):
        """Call libRDF roqet utility"""
        pass

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('-o', '--output', type=str, default=None,
              help="Specifiy variable to hold output")
    @argument('-s', '--source', default=None,
              help="Source can be a literal, a python string, or python list")
    @argument('-c', '--count', default=False, action="store_true",
              help="output count of rows returned")
    @cell_magic
    def sparql(self, line, cell=None):
        arg = parse_argstring(self.sparql, line)

        if cell is not None:
            sources, cell = extract_froms(cell)

        sources.extend(self._parse_source(arg.source))

        if arg.model is None and len(sources) == 0:
            raise UsageError("Please specify a source to query against.")

        model = self._get_model(arg.model)

        for source in sources:
            if source.startswith("tracker:"):
                raise NotImplemented("Tracker queries not implemented yet")
            else:
                load_source(model, source)

        body = prepare_query(cell)
        query = RDF.SPARQLQuery(body)
        results = LibRdfResults(query.execute(model))

        if arg.count:
            print ("Found {0} rows.".format(len(results)))

        if arg.output is None:
            return results
        else:
            self.shell.user_ns[arg.output] = results

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('sources', nargs='*', default=None,
              help="Source can be a literal, a python string, or python list")
    @line_magic
    def load_source(self, line):
        arg = parse_argstring(self.load_source, line)
        model = self._get_model(arg.model)

        sources = []
        if arg.sources is not None:
            for s in arg.sources:
                sources.extend(self._parse_source(s))

        for s in sources:
            load_source(model, s)

        if arg.model is not None:
            return model

    def _get_model(self, variable):
        """Get model from user name space, or make a new one"""
        if variable is not None:
            m = self.shell.user_ns.get(variable)
            if not isinstance(m, RDF.Model):
                raise ValueError("Model needs to be RDF.Model")
        else:
            m = make_temp_model()
        return m

    def _parse_source(self, variable):
        """parse a source argument and return a list of sources
        """
        sources = []
        if variable is not None:
            source = self.shell.user_ns.get(variable, variable)
            if type(source) in types.StringTypes:
                sources.append(source)
            elif isinstance(source, collections.Iterable):
                sources.extend(source)
        return sources

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('-f', '--format', default='turtle',
              help="Specify format to save model as")
    @argument('filename', nargs=1, type=str, help="filename to save model as")
    @line_magic
    def save_model(self, line):
        arg = parse_argstring(self.load_source, line)
        model = self._get_model(arg.model)
        if model is None:
            UsageError("Please specify a model to save")

        if arg.filename is None:
            UsageError("Please specify a destination")

        serializer = RDF.Serializer(name=arg.name)
        with open(arg.filename, 'w') as outstream:
            outstream.write(serializer.serialize_model_to_string())
            outstream.write('\n')

def extract_froms(cell, remove=True):
    """Extract from statements from sparql query
    librdf doesn't actually use the froms.

    if remove is true, remove the from statements from the query
    returns ([url, url], query)
    """
    from_re = re.compile(
        "from <(?P<url>[A-Za-z0-9!@#$%^&*()-_\\|\'\"/?\]\[{}+=]+)>",
        re.IGNORECASE)
    froms = []
    match = from_re.search(cell)
    while match is not None:
        froms.append(match.group('url'))
        if remove:
            cell = from_re.sub('', cell, 1)
        match = from_re.search(cell)
    return froms, cell

def prepare_query(cell):
    """Add additional our namespaces to head of query
    """
    query = []
    template = "PREFIX {prefix}: <{url}>"
    for prefix, namespace in _prefixes.items():
        query.append(template.format(prefix=prefix, url=namespace))
    query.append(cell)
    return os.linesep.join(query)

def make_temp_model():
    """Make a scratch model
    """
    s = RDF.MemoryStorage()
    m = RDF.Model(s)
    return m

def load_source(model, source):
    if isinstance(source, RDF.Node):
        source = str(source.uri)

    url = urlparse(source)
    if url.scheme is None:
        source = 'file://'+os.path.abspath(source)
        url = urlparse(source)
    if url.scheme in ('http', 'https'):
        # remote,
        stream = urlopen(source)
        content_type = stream.headers.get('content-type')
        if stream.code == 200:
            parser = guess_parser(content_type, url.path)
        else:
            raise HTTPError("Problem opening {}: {}".format(
                source, stream.code))

    elif url.scheme in ('file'):
        # local
        if not os.path.exists(url.path):
            raise IOError("File %s does not exist" % (url.path,))
        stream = open(url.path, 'r')
        parser = guess_parser(None, url.path)

    body = stream.read()
    stream.close()
    parser.parse_string_into_model(model, body, source)

def guess_parser(content_type, pathname):
    name = guess_parser_name(content_type, pathname)
    return RDF.Parser(name=name)

def guess_parser_name(content_type, pathname):
    if content_type is None or content_type.startswith('text/plain'):
        return guess_parser_name_by_extension(pathname)
    if content_type.startswith('application/rdf+xml'):
        return 'rdfxml'
    elif content_type.startswith('application/x-turtle'):
        return 'turtle'
    elif content_type.startswith('text/turtle'):
        return 'turtle'
    elif content_type.startswith('text/html'):
        return 'rdfa'
    else:
        return 'guess'

def guess_parser_name_by_extension(pathname):
    _, ext = os.path.splitext(pathname)
    if ext in ('.xml', '.rdf'):
        return 'rdfxml'
    elif ext in ('.html','.xhtml'):
        return 'rdfa'
    elif ext in ('.turtle','.ttl'):
        return 'turtle'
    return 'guess'

_loaded = False
def load_ipython_extension(ip):
    """Load the extension in IPython."""
    global _loaded
    if not _loaded:
        ip.register_magics(SPARQLMagics)
        _loaded = True
<commit_msg>update parse_sources to handle reading sources out of a LibRdfResults object<commit_after># Copyright (C) 2014 Diane Trout
#
# This package is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License Version 2.1
# as published by the Free Software Foundation or any newer version.

# This package is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
# License Version 2.1 for more details.

# You should have received a copy of the GNU Lesser General Public
# License Version 2.1 along with this package; if not, write to the
# Free Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston,
# MA 02110-1301 USA

import os
import re
import six

from six.moves import range
from six.moves.urllib.error import HTTPError
from six.moves.urllib.parse import urlparse
from six.moves.urllib.request import urlopen

import RDF
import collections

from IPython.core.magic import (Magics, magics_class, cell_magic, line_magic,
                                line_cell_magic, UsageError)
from IPython.core.magic_arguments import (
    argument, magic_arguments, parse_argstring)

_prefixes = {
    'rdf': "http://www.w3.org/1999/02/22-rdf-syntax-ns#",
    'rdfs': "http://www.w3.org/2000/01/rdf-schema#",
    'owl': 'http://www.w3.org/2002/07/owl#',
    'dc': "http://purl.org/dc/elements/1.1/",
    'xml': 'http://www.w3.org/XML/1998/namespace',
    'xsd': "http://www.w3.org/2001/XMLSchema#",
    'vs': 'http://www.w3.org/2003/06/sw-vocab-status/ns#',
    'wot': 'http://xmlns.com/wot/0.1/'
}
_namespaces = dict(( (v, k) for k, v in _prefixes.items()))

def dump_model(model):
    export = RDF.Serializer(name='turtle')
    print (export.serialize_model_to_string(model))

def display_node(node, mime_type):
    if mime_type == 'text/html':
        template = '{curie}'
    else:
        template = '{curie}'

    if isinstance(node, RDF.Node) and node.is_resource():
        match = ""
        uri = str(node.uri)
        curie = uri
        for namespace, prefix in _namespaces.items():
            if uri.startswith(namespace) and len(namespace) > len(match):
                match = namespace

        if len(match) > 0:
            prefix = _namespaces[match]
            curie = uri.replace(match, "{0}:".format(prefix))

        return template.format(anchor=uri, curie=curie)
    else:
        return str(node)

class Node(RDF.Node):
    def _repr_html_(self):
        return display_node(self, mime_type='text/html')

Uri = RDF.Uri
NS = RDF.NS

class LibRdfResults(collections.Sequence):
    results = None
    columns = None
    def __init__(self, result_set):
        if result_set is not None:
            self.columns = self.get_bindings(result_set)
            self.results = self.get_results(result_set)

    def get_bindings(self, result_set):
        b = []
        for i in range(result_set.get_bindings_count()):
            b.append(result_set.get_binding_name(i))
        return b

    def get_results(self, result_set):
        results = []
        for row in result_set:
            record = [ row[k] for k in self.columns ]
            results.append(record)
        return results

    def __getitem__(self, key):
        return dict(zip(self.columns, self.results[key]))

    def __len__(self):
        return len(self.results)

    def generate_html(self):
        result_set = iter(self)
        yield '

'
        for c in self.columns:
            yield "{0}	".format(str(c))
        yield '
'
        for row in self:
            yield '
'
            for key in self.columns:
                value = row[key]
                yield ''
                yield display_node(value, mime_type='text/html')
                yield '	'
            yield '
'
        yield '
'

    def _repr_html_(self):
        return "".join(self.generate_html())

    def __str__(self):
        result_set = iter(self)
        output = []
        output.append("\t".join(self.columns))
        for row in self.results:
            ordered = (display_node(value, mime_type='text/plain')
                       for value in row)
            output.append("\t".join(ordered))
        return os.linesep.join(output)

@magics_class
class SPARQLMagics(Magics):
    SUPPORT_FROM = False
    def __init__(self, shell):
        super(SPARQLMagics, self).__init__(shell)

    @magic_arguments()
    @argument('prefix', nargs=1, type=str,
              help="Specify namespace prefix")
    @argument('namespace', nargs=1, type=str,
              help="specify namespace")
    @line_magic
    def addns(self, line):
        args = parse_argstring(self.addns, line)
        if args.prefix is None or args.namespace is None:
            raise UsageError("Please specify prefix and namespace")
        namespace = args.namespace[0]
        prefix = args.prefix[0]
        _namespaces[namespace] = prefix
        _prefixes[prefix] = namespace
        self.shell.user_ns[prefix] = RDF.NS(namespace)

    @line_magic
    def lsns(self, line=None):
        keys = sorted(_prefixes.keys())
        maxlen = max((len(k) for k in keys))
        for k in keys:
            print ("{prefix:{width}} {ns}".format(
                prefix=k,
                width=maxlen,
                ns=_prefixes[k]))

    @magic_arguments()
    @argument('name', nargs=1, type=str,
              help="remove an item from our default namespaces")
    @line_magic
    def delns(self, line=None):
        args = parse_argstring(self.delns, line)
        if args.name is None:
            raise UsageError("Specify what to delete")

        name = args.name[0]
        if name in _prefixes:
            del _namespaces[self._prefixes[name]]
            del _prefixes[name]
            del self.shell.user_ns[name]
        elif name in _namespaces:
            prefix = self._namespaces[name]
            del _prefixes[prefix]
            del _namespaces[name]
            del self.shell.user_ns[prefix]
        else:
            raise UsageError("%s was not found in our namespace cache" % (
                name,))

    @magic_arguments()
    @argument('-D', '--data', default=None, type=str,
              help="Add RDF Data source URI")
    @argument('-p', '--protocol', type=str, default=None,
              help="Call a SPARQL HTTP Protocol at service URI")
    @argument('-m', '--model', type=str, default=None,
              help="Specify a model to use")
    @argument('--output', type=str,
              default=None,
              help="Specifiy variable to hold output")
    @line_cell_magic
    def roqet(self, line, cell):
        """Call libRDF roqet utility"""
        pass

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('-o', '--output', type=str, default=None,
              help="Specifiy variable to hold output")
    @argument('-s', '--source', default=None,
              help="Source can be a literal, a python string, or python list")
    @argument('-c', '--count', default=False, action="store_true",
              help="output count of rows returned")
    @cell_magic
    def sparql(self, line, cell=None):
        arg = parse_argstring(self.sparql, line)

        if cell is not None:
            sources, cell = extract_froms(cell)

        source = self.shell.user_ns.get(arg.source, arg.source)
        sources.extend(self._parse_source(source))

        if arg.model is None and len(sources) == 0:
            raise UsageError("Please specify a source to query against.")

        model = self._get_model(arg.model)

        for source in sources:
            if source.startswith("tracker:"):
                raise NotImplemented("Tracker queries not implemented yet")
            else:
                load_source(model, source)

        body = prepare_query(cell)
        query = RDF.SPARQLQuery(body)
        results = LibRdfResults(query.execute(model))

        if arg.count:
            print ("Found {0} rows.".format(len(results)))

        if arg.output is None:
            return results
        else:
            self.shell.user_ns[arg.output] = results

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('sources', nargs='*', default=None,
              help="Source can be a literal, a python string, or python list")
    @line_magic
    def load_source(self, line):
        arg = parse_argstring(self.load_source, line)
        model = self._get_model(arg.model)

        sources = []
        if arg.sources is not None:
            for variable_name in arg.sources:
                source = self.shell.user_ns.get(variable_name, variable_name)
                sources.extend(self._parse_source(source))

        for s in sources:
            load_source(model, s)

        if arg.model is not None:
            return model

    def _get_model(self, variable):
        """Get model from user name space, or make a new one"""
        if variable is not None:
            m = self.shell.user_ns.get(variable)
            if not isinstance(m, RDF.Model):
                raise ValueError("Model needs to be RDF.Model")
        else:
            m = make_temp_model()
        return m

    def _parse_source(self, source):
        """parse a source argument and return a list of sources
        """
        sources = []
        if isinstance(source, six.string_types):
            sources.append(source)
        elif isinstance(source, RDF.Node):
            sources.append(source)
        elif isinstance(source, LibRdfResults):
            for result in source:
                for value in result.values():
                    if isinstance(value, RDF.Node):
                        sources.append(value)
        elif isinstance(source, collections.Iterable):
            sources.extend(source)
        return sources

    @magic_arguments()
    @argument('-m', '--model', default=None,
              help="use specified variable as the model to store "\
                   "intermediate results in.")
    @argument('-f', '--format', default='turtle',
              help="Specify format to save model as")
    @argument('filename', nargs=1, type=str, help="filename to save model as")
    @line_magic
    def save_model(self, line):
        arg = parse_argstring(self.load_source, line)
        model = self._get_model(arg.model)
        if model is None:
            UsageError("Please specify a model to save")

        if arg.filename is None:
            UsageError("Please specify a destination")

        serializer = RDF.Serializer(name=arg.name)
        with open(arg.filename, 'w') as outstream:
            outstream.write(serializer.serialize_model_to_string())
            outstream.write('\n')

def extract_froms(cell, remove=True):
    """Extract from statements from sparql query
    librdf doesn't actually use the froms.

    if remove is true, remove the from statements from the query
    returns ([url, url], query)
    """
    from_re = re.compile(
        "from <(?P<url>[A-Za-z0-9!@#$%^&*()-_\\|\'\"/?\]\[{}+=]+)>",
        re.IGNORECASE)
    froms = []
    match = from_re.search(cell)
    while match is not None:
        froms.append(match.group('url'))
        if remove:
            cell = from_re.sub('', cell, 1)
        match = from_re.search(cell)
    return froms, cell

def prepare_query(cell):
    """Add additional our namespaces to head of query
    """
    query = []
    template = "PREFIX {prefix}: <{url}>"
    for prefix, namespace in _prefixes.items():
        query.append(template.format(prefix=prefix, url=namespace))
    query.append(cell)
    return os.linesep.join(query)

def make_temp_model():
    """Make a scratch model
    """
    s = RDF.MemoryStorage()
    m = RDF.Model(s)
    return m

def load_source(model, source):
    if isinstance(source, RDF.Node):
        source = str(source.uri)

    url = urlparse(source)
    if url.scheme is None:
        source = 'file://'+os.path.abspath(source)
        url = urlparse(source)
    if url.scheme in ('http', 'https'):
        # remote,
        stream = urlopen(source)
        content_type = stream.headers.get('content-type')
        if stream.code == 200:
            parser = guess_parser(content_type, url.path)
        else:
            raise HTTPError("Problem opening {}: {}".format(
                source, stream.code))

    elif url.scheme in ('file'):
        # local
        if not os.path.exists(url.path):
            raise IOError("File %s does not exist" % (url.path,))
        stream = open(url.path, 'r')
        parser = guess_parser(None, url.path)

    body = stream.read()
    stream.close()
    parser.parse_string_into_model(model, body, source)

def guess_parser(content_type, pathname):
    name = guess_parser_name(content_type, pathname)
    return RDF.Parser(name=name)

def guess_parser_name(content_type, pathname):
    if content_type is None or content_type.startswith('text/plain'):
        return guess_parser_name_by_extension(pathname)
    if content_type.startswith('application/rdf+xml'):
        return 'rdfxml'
    elif content_type.startswith('application/x-turtle'):
        return 'turtle'
    elif content_type.startswith('text/turtle'):
        return 'turtle'
    elif content_type.startswith('text/html'):
        return 'rdfa'
    else:
        return 'guess'

def guess_parser_name_by_extension(pathname):
    _, ext = os.path.splitext(pathname)
    if ext in ('.xml', '.rdf'):
        return 'rdfxml'
    elif ext in ('.html','.xhtml'):
        return 'rdfa'
    elif ext in ('.turtle','.ttl'):
        return 'turtle'
    return 'guess'

_loaded = False
def load_ipython_extension(ip):
    """Load the extension in IPython."""
    global _loaded
    if not _loaded:
        ip.register_magics(SPARQLMagics)
        _loaded = True
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import division

import numpy as np
from numpy.testing import assert_equal
from ...tests.helper import pytest
from ...tests.compat import assert_allclose

from .. import funcs
from ...utils.misc import NumpyRNGContext

try:
    from scipy import stats  # used in testing
except ImportError:
    HAS_SCIPY = False
else:
    HAS_SCIPY = True

def test_sigma_clip():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):
        # Amazing, I've got the same combination on my luggage!
        randvar = randn(10000)

        filtered_data = funcs.sigma_clip(randvar, 1, 2)

        assert sum(filtered_data.mask) > 0
        assert sum(~filtered_data.mask) < randvar.size

        #this is actually a silly thing to do, because it uses the standard
        #deviation as the variance, but it tests to make sure these arguments
        #are actually doing something
        filtered_data2 = funcs.sigma_clip(randvar, 1, 2, varfunc=np.std)
        assert not np.all(filtered_data.mask == filtered_data2.mask)

        filtered_data3 = funcs.sigma_clip(randvar, 1, 2, cenfunc=np.mean)
        assert not np.all(filtered_data.mask == filtered_data3.mask)

        # make sure the iters=None method works at all.
        filtered_data = funcs.sigma_clip(randvar, 3, None)

        # test copying
        assert filtered_data.data[0] == randvar[0]
        filtered_data.data[0] += 1.
        assert filtered_data.data[0] != randvar[0]

        filtered_data = funcs.sigma_clip(randvar, 3, None, copy=False)
        assert filtered_data.data[0] == randvar[0]
        filtered_data.data[0] += 1.
        assert filtered_data.data[0] == randvar[0]

        # test axis
        data = np.arange(5)+np.random.normal(0.,0.05,(5,5))+np.diag(np.ones(5))
        filtered_data = funcs.sigma_clip(data, axis=0, sig=2.3)
        assert filtered_data.count() == 20
        filtered_data = funcs.sigma_clip(data, axis=1, sig=2.3)
        assert filtered_data.count() == 25

def test_median_absolute_deviation():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        mad = funcs.median_absolute_deviation(randvar)

        #test whether an array is returned if an axis is used
        randvar = randvar.reshape((10, 1000))
        mad = funcs.median_absolute_deviation(randvar, axis=1)
        assert len(mad) == 10
        assert mad.size < randvar.size

def test_biweight_location():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        cbl = funcs.biweight_location(randvar)

        assert abs(cbl-0) < 1e-2

def test_biweight_midvariance():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        scl = funcs.biweight_midvariance(randvar)

        assert abs(scl-1) < 1e-2

@pytest.mark.skipif('not HAS_SCIPY')
def test_compare_to_scipy_sigmaclip():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        randvar = randn(10000)

        astropyres = funcs.sigma_clip(randvar, 3, None, np.mean)
        scipyres = stats.sigmaclip(randvar, 3, 3)[0]

        assert astropyres.count() == len(scipyres)
        assert_equal(astropyres[~astropyres.mask].data, scipyres)

@pytest.mark.skipif('not HAS_SCIPY')
def test_binom_conf_interval():

    # Test Wilson and Jeffreys interval for corner cases:
    # Corner cases: k = 0, k = n, conf = 0., conf = 1.
    n = 5
    k = [0, 4, 5]
    for conf in [0., 0.5, 1.]:
        res = funcs.binom_conf_interval(k, n, conf=conf, interval='wilson')
        assert ((res >= 0.) & (res <= 1.)).all()
        res = funcs.binom_conf_interval(k, n, conf=conf, interval='jeffreys')
        assert ((res >= 0.) & (res <= 1.)).all()

    # Test Jeffreys interval accuracy against table in Brown et al. (2001).
    # (See `binom_conf_interval` docstring for reference.)
    k = [0, 1, 2, 3, 4]
    n = 7
    conf = 0.95
    result = funcs.binom_conf_interval(k, n, conf=conf, interval='jeffreys')
    table = np.array([[0.000, 0.016, 0.065, 0.139, 0.234],
                      [0.292, 0.501, 0.648, 0.766, 0.861]])
    assert_allclose(result, table, atol=1.e-3, rtol=0.)

    # Test Wald interval
    result = funcs.binom_conf_interval(0, 5, interval='wald')
    assert_allclose(result, 0.)  # conf interval is [0, 0] when k = 0
    result = funcs.binom_conf_interval(5, 5, interval='wald')
    assert_allclose(result, 1.)  # conf interval is [1, 1] when k = n
    result = funcs.binom_conf_interval(500, 1000, conf=0.68269,
                                       interval='wald')
    assert_allclose(result[0], 0.5 - 0.5 / np.sqrt(1000.))
    assert_allclose(result[1], 0.5 + 0.5 / np.sqrt(1000.))

@pytest.mark.skipif('not HAS_SCIPY')
def test_binned_binom_proportion():

    # Check that it works.
    nbins = 20
    x = np.linspace(0., 10., 100)  # Guarantee an `x` in every bin.
    success = np.ones(len(x), dtype=np.bool)
    bin_ctr, bin_hw, p, perr = funcs.binned_binom_proportion(x, success,
                                                             bins=nbins)

    # Check shape of outputs
    assert bin_ctr.shape == (nbins,)
    assert bin_hw.shape == (nbins,)
    assert p.shape == (nbins,)
    assert perr.shape == (2, nbins)

    # Check that p is 1 in all bins, since success = True for all `x`.
    assert (p == 1.).all()

    # Check that p is 0 in all bins if success = False for all `x`.
    success[:] = False
    bin_ctr, bin_hw, p, perr = funcs.binned_binom_proportion(x, success,
                                                             bins=nbins)
    assert (p == 0.).all()

def test_signal_to_noise_oir_ccd():

    result = funcs.signal_to_noise_oir_ccd(1,25,0,0,0,1)
    assert 5.0 == result
    #check to make sure gain works
    result = funcs.signal_to_noise_oir_ccd(1,5,0,0,0,1,5)
    assert 5.0 == result

    #now add in sky, dark current, and read noise
    #make sure the snr goes down
    result = funcs.signal_to_noise_oir_ccd(1,25,1,0,0,1)
    assert result < 5.0
    result = funcs.signal_to_noise_oir_ccd(1,25,0,1,0,1)
    assert result < 5.0
    result = funcs.signal_to_noise_oir_ccd(1,25,0,0,1,1)
    assert result < 5.0

    #make sure snr increases with time
    result = funcs.signal_to_noise_oir_ccd(2,25,0,0,0,1)
    assert result > 5.0

def test_bootstrap():
    bootarr = np.array([1,2,3,4,5,6,7,8,9,0])
    #test general bootstrapping
    answer = np.array([[7,4,8,5,7,0,3,7,8,5],[4,8,8,3,6,5,2,8,6,2]])
    with NumpyRNGContext(42):
        assert answer = funcs.bootstrap(bootarr,2)

    #test with a bootfunction
    with NumpyRNGContext(42):
        bootresult = np.mean(bootstrap(a,10000,bootfunc=np.mean))
        assert_allclose(np.mean(bootarr),bootresult,atol=0.01)

<commit_msg>fixed silly mistake in tests<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import division

import numpy as np
from numpy.testing import assert_equal
from ...tests.helper import pytest
from ...tests.compat import assert_allclose

from .. import funcs
from ...utils.misc import NumpyRNGContext

try:
    from scipy import stats  # used in testing
except ImportError:
    HAS_SCIPY = False
else:
    HAS_SCIPY = True

def test_sigma_clip():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):
        # Amazing, I've got the same combination on my luggage!
        randvar = randn(10000)

        filtered_data = funcs.sigma_clip(randvar, 1, 2)

        assert sum(filtered_data.mask) > 0
        assert sum(~filtered_data.mask) < randvar.size

        #this is actually a silly thing to do, because it uses the standard
        #deviation as the variance, but it tests to make sure these arguments
        #are actually doing something
        filtered_data2 = funcs.sigma_clip(randvar, 1, 2, varfunc=np.std)
        assert not np.all(filtered_data.mask == filtered_data2.mask)

        filtered_data3 = funcs.sigma_clip(randvar, 1, 2, cenfunc=np.mean)
        assert not np.all(filtered_data.mask == filtered_data3.mask)

        # make sure the iters=None method works at all.
        filtered_data = funcs.sigma_clip(randvar, 3, None)

        # test copying
        assert filtered_data.data[0] == randvar[0]
        filtered_data.data[0] += 1.
        assert filtered_data.data[0] != randvar[0]

        filtered_data = funcs.sigma_clip(randvar, 3, None, copy=False)
        assert filtered_data.data[0] == randvar[0]
        filtered_data.data[0] += 1.
        assert filtered_data.data[0] == randvar[0]

        # test axis
        data = np.arange(5)+np.random.normal(0.,0.05,(5,5))+np.diag(np.ones(5))
        filtered_data = funcs.sigma_clip(data, axis=0, sig=2.3)
        assert filtered_data.count() == 20
        filtered_data = funcs.sigma_clip(data, axis=1, sig=2.3)
        assert filtered_data.count() == 25

def test_median_absolute_deviation():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        mad = funcs.median_absolute_deviation(randvar)

        #test whether an array is returned if an axis is used
        randvar = randvar.reshape((10, 1000))
        mad = funcs.median_absolute_deviation(randvar, axis=1)
        assert len(mad) == 10
        assert mad.size < randvar.size

def test_biweight_location():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        cbl = funcs.biweight_location(randvar)

        assert abs(cbl-0) < 1e-2

def test_biweight_midvariance():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        #test that it runs
        randvar = randn(10000)
        scl = funcs.biweight_midvariance(randvar)

        assert abs(scl-1) < 1e-2

@pytest.mark.skipif('not HAS_SCIPY')
def test_compare_to_scipy_sigmaclip():
    from numpy.random import randn

    #need to seed the numpy RNG to make sure we don't get some amazingly flukey
    #random number that breaks one of the tests

    with NumpyRNGContext(12345):

        randvar = randn(10000)

        astropyres = funcs.sigma_clip(randvar, 3, None, np.mean)
        scipyres = stats.sigmaclip(randvar, 3, 3)[0]

        assert astropyres.count() == len(scipyres)
        assert_equal(astropyres[~astropyres.mask].data, scipyres)

@pytest.mark.skipif('not HAS_SCIPY')
def test_binom_conf_interval():

    # Test Wilson and Jeffreys interval for corner cases:
    # Corner cases: k = 0, k = n, conf = 0., conf = 1.
    n = 5
    k = [0, 4, 5]
    for conf in [0., 0.5, 1.]:
        res = funcs.binom_conf_interval(k, n, conf=conf, interval='wilson')
        assert ((res >= 0.) & (res <= 1.)).all()
        res = funcs.binom_conf_interval(k, n, conf=conf, interval='jeffreys')
        assert ((res >= 0.) & (res <= 1.)).all()

    # Test Jeffreys interval accuracy against table in Brown et al. (2001).
    # (See `binom_conf_interval` docstring for reference.)
    k = [0, 1, 2, 3, 4]
    n = 7
    conf = 0.95
    result = funcs.binom_conf_interval(k, n, conf=conf, interval='jeffreys')
    table = np.array([[0.000, 0.016, 0.065, 0.139, 0.234],
                      [0.292, 0.501, 0.648, 0.766, 0.861]])
    assert_allclose(result, table, atol=1.e-3, rtol=0.)

    # Test Wald interval
    result = funcs.binom_conf_interval(0, 5, interval='wald')
    assert_allclose(result, 0.)  # conf interval is [0, 0] when k = 0
    result = funcs.binom_conf_interval(5, 5, interval='wald')
    assert_allclose(result, 1.)  # conf interval is [1, 1] when k = n
    result = funcs.binom_conf_interval(500, 1000, conf=0.68269,
                                       interval='wald')
    assert_allclose(result[0], 0.5 - 0.5 / np.sqrt(1000.))
    assert_allclose(result[1], 0.5 + 0.5 / np.sqrt(1000.))

@pytest.mark.skipif('not HAS_SCIPY')
def test_binned_binom_proportion():

    # Check that it works.
    nbins = 20
    x = np.linspace(0., 10., 100)  # Guarantee an `x` in every bin.
    success = np.ones(len(x), dtype=np.bool)
    bin_ctr, bin_hw, p, perr = funcs.binned_binom_proportion(x, success,
                                                             bins=nbins)

    # Check shape of outputs
    assert bin_ctr.shape == (nbins,)
    assert bin_hw.shape == (nbins,)
    assert p.shape == (nbins,)
    assert perr.shape == (2, nbins)

    # Check that p is 1 in all bins, since success = True for all `x`.
    assert (p == 1.).all()

    # Check that p is 0 in all bins if success = False for all `x`.
    success[:] = False
    bin_ctr, bin_hw, p, perr = funcs.binned_binom_proportion(x, success,
                                                             bins=nbins)
    assert (p == 0.).all()

def test_signal_to_noise_oir_ccd():

    result = funcs.signal_to_noise_oir_ccd(1,25,0,0,0,1)
    assert 5.0 == result
    #check to make sure gain works
    result = funcs.signal_to_noise_oir_ccd(1,5,0,0,0,1,5)
    assert 5.0 == result

    #now add in sky, dark current, and read noise
    #make sure the snr goes down
    result = funcs.signal_to_noise_oir_ccd(1,25,1,0,0,1)
    assert result < 5.0
    result = funcs.signal_to_noise_oir_ccd(1,25,0,1,0,1)
    assert result < 5.0
    result = funcs.signal_to_noise_oir_ccd(1,25,0,0,1,1)
    assert result < 5.0

    #make sure snr increases with time
    result = funcs.signal_to_noise_oir_ccd(2,25,0,0,0,1)
    assert result > 5.0

def test_bootstrap():
    bootarr = np.array([1,2,3,4,5,6,7,8,9,0])
    #test general bootstrapping
    answer = np.array([[7,4,8,5,7,0,3,7,8,5],[4,8,8,3,6,5,2,8,6,2]])
    with NumpyRNGContext(42):
        assert answer == funcs.bootstrap(bootarr,2)

    #test with a bootfunction
    with NumpyRNGContext(42):
        bootresult = np.mean(bootstrap(a,10000,bootfunc=np.mean))
        assert_allclose(np.mean(bootarr),bootresult,atol=0.01)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
    :codeauthor: Pedro Algarvio (pedro@algarvio.me)

    salt.utils.odict
    ~~~~~~~~~~~~~~~~

    This is a compatibility/"importability" layer for an ordered dictionary.
    Tries to import from the standard library if python >= 2.7, then from the
    ``ordereddict`` package available from PyPi, and, as a last resort,
    provides an ``OrderedDict`` implementation based on::

        http://code.activestate.com/recipes/576669/

    It also implements a DefaultOrderedDict Class that serves  as a
    combination of ``OrderedDict`` and ``defaultdict``
    It's source was submitted here::

        http://stackoverflow.com/questions/6190331/
'''

# Import python libs
from __future__ import absolute_import, unicode_literals, print_function
from collections import Callable

# Import 3rd-party libs
from salt.ext import six

try:
    # pylint: disable=E0611,minimum-python-version
    import collections

    class OrderedDict(collections.OrderedDict):
        __hash__ = None
except (ImportError, AttributeError):
    try:
        import ordereddict

        class OrderedDict(ordereddict.OrderedDict):  # pylint: disable=W0232
            __hash_ = None
    except ImportError:
        # {{{ http://code.activestate.com/recipes/576693/ (r9)
        # Backport of OrderedDict() class that runs on Python 2.4, 2.5, 2.6, 2.7 and pypy.
        # Passes Python2.7's test suite and incorporates all the latest updates.

        # pylint: disable=import-error,no-name-in-module
        try:
            from salt.ext.six.moves._thread import get_ident as _get_ident
        except ImportError:
            from salt.ext.six.moves._dummy_thread import get_ident as _get_ident
        # pylint: enable=import-error,no-name-in-module

#        try:
#            from _abcoll import KeysView, ValuesView, ItemsView
#        except ImportError:
#            pass

        class OrderedDict(dict):
            'Dictionary that remembers insertion order'
            # An inherited dict maps keys to values.
            # The inherited dict provides __getitem__, __len__, __contains__, and get.
            # The remaining methods are order-aware.
            # Big-O running times for all methods are the same as for regular dictionaries.

            # The internal self.__map dictionary maps keys to links in a doubly linked list.
            # The circular doubly linked list starts and ends with a sentinel element.
            # The sentinel element never gets deleted (this simplifies the algorithm).
            # Each link is stored as a list of length three:  [PREV, NEXT, KEY].
            __hash_ = None

            def __init__(self, *args, **kwds):  # pylint: disable=E1003
                '''Initialize an ordered dictionary.  Signature is the same as for
                regular dictionaries, but keyword arguments are not recommended
                because their insertion order is arbitrary.

                '''
                super(OrderedDict, self).__init__()  # pylint: disable=E1003
                if len(args) > 1:
                    raise TypeError(
                        'expected at most 1 arguments, got {0}'.format(len(args))
                    )
                try:
                    self.__root
                except AttributeError:
                    self.__root = root = []                     # sentinel node
                    root[:] = [root, root, None]
                    self.__map = {}
                self.__update(*args, **kwds)

            def __setitem__(self, key, value, dict_setitem=dict.__setitem__):
                'od.__setitem__(i, y) <==> od[i]=y'
                # Setting a new item creates a new link which goes at the end of the linked
                # list, and the inherited dictionary is updated with the new key/value pair.
                if key not in self:
                    root = self.__root
                    last = root[0]
                    last[1] = root[0] = self.__map[key] = [last, root, key]
                dict_setitem(self, key, value)

            def __delitem__(self, key, dict_delitem=dict.__delitem__):
                'od.__delitem__(y) <==> del od[y]'
                # Deleting an existing item uses self.__map to find the link which is
                # then removed by updating the links in the predecessor and successor nodes.
                dict_delitem(self, key)
                link_prev, link_next, key = self.__map.pop(key)
                link_prev[1] = link_next
                link_next[0] = link_prev

            def __iter__(self):
                'od.__iter__() <==> iter(od)'
                root = self.__root
                curr = root[1]
                while curr is not root:
                    yield curr[2]
                    curr = curr[1]

            def __reversed__(self):
                'od.__reversed__() <==> reversed(od)'
                root = self.__root
                curr = root[0]
                while curr is not root:
                    yield curr[2]
                    curr = curr[0]

            def clear(self):
                'od.clear() -> None.  Remove all items from od.'
                try:
                    for node in six.itervalues(self.__map):
                        del node[:]
                    root = self.__root
                    root[:] = [root, root, None]
                    self.__map.clear()
                except AttributeError:
                    pass
                dict.clear(self)

            def popitem(self, last=True):
                '''od.popitem() -> (k, v), return and remove a (key, value) pair.
                Pairs are returned in LIFO order if last is true or FIFO order if false.

                '''
                if not self:
                    raise KeyError('dictionary is empty')
                root = self.__root
                if last:
                    link = root[0]
                    link_prev = link[0]
                    link_prev[1] = root
                    root[0] = link_prev
                else:
                    link = root[1]
                    link_next = link[1]
                    root[1] = link_next
                    link_next[0] = root
                key = link[2]
                del self.__map[key]
                value = dict.pop(self, key)
                return key, value

            # -- the following methods do not depend on the internal structure --

            def keys(self):
                'od.keys() -> list of keys in od'
                return list(self)

            def values(self):
                'od.values() -> list of values in od'
                return [self[key] for key in self]

            def items(self):
                'od.items() -> list of (key, value) pairs in od'
                return [(key, self[key]) for key in self]

            def iterkeys(self):
                'od.iterkeys() -> an iterator over the keys in od'
                return iter(self)

            def itervalues(self):
                'od.itervalues -> an iterator over the values in od'
                for k in self:
                    yield self[k]

            def iteritems(self):
                'od.iteritems -> an iterator over the (key, value) items in od'
                for k in self:
                    yield (k, self[k])

            def update(*args, **kwds):  # pylint: disable=E0211
                '''od.update(E, **F) -> None.  Update od from dict/iterable E and F.

                If E is a dict instance, does:           for k in E: od[k] = E[k]
                If E has a .keys() method, does:         for k in E.keys(): od[k] = E[k]
                Or if E is an iterable of items, does:   for k, v in E: od[k] = v
                In either case, this is followed by:     for k, v in F.items(): od[k] = v

                '''
                if len(args) > 2:
                    raise TypeError(
                        'update() takes at most 2 positional '
                        'arguments ({0} given)'.format(len(args))
                    )
                elif not args:
                    raise TypeError('update() takes at least 1 argument (0 given)')
                self = args[0]
                # Make progressively weaker assumptions about "other"
                other = ()
                if len(args) == 2:
                    other = args[1]
                if isinstance(other, dict):
                    for key in other:
                        self[key] = other[key]
                elif hasattr(other, 'keys'):
                    for key in other:
                        self[key] = other[key]
                else:
                    for key, value in other:
                        self[key] = value
                for key, value in six.iteritems(kwds):
                    self[key] = value

            __update = update  # let subclasses override update without breaking __init__

            __marker = object()

            def pop(self, key, default=__marker):
                '''od.pop(k[,d]) -> v, remove specified key and return the corresponding value.
                If key is not found, d is returned if given, otherwise KeyError is raised.

                '''
                if key in self:
                    result = self[key]
                    del self[key]
                    return result
                if default is self.__marker:
                    raise KeyError(key)
                return default

            def setdefault(self, key, default=None):
                'od.setdefault(k[,d]) -> od.get(k,d), also set od[k]=d if k not in od'
                if key in self:
                    return self[key]
                self[key] = default
                return default

            def __repr__(self, _repr_running={}):  # pylint: disable=W0102
                'od.__repr__() <==> repr(od)'
                call_key = id(self), _get_ident()
                if call_key in _repr_running:
                    return '...'
                _repr_running[call_key] = 1
                try:
                    if not self:
                        return '{0}()'.format(self.__class__.__name__)
                    return '{0}(\'{1}\')'.format(self.__class__.__name__, list(self.items()))
                finally:
                    del _repr_running[call_key]

            def __reduce__(self):
                'Return state information for pickling'
                items = [[k, self[k]] for k in self]
                inst_dict = vars(self).copy()
                for k in vars(OrderedDict()):
                    inst_dict.pop(k, None)
                if inst_dict:
                    return (self.__class__, (items,), inst_dict)
                return self.__class__, (items,)

            def copy(self):
                'od.copy() -> a shallow copy of od'
                return self.__class__(self)

            @classmethod
            def fromkeys(cls, iterable, value=None):
                '''OD.fromkeys(S[, v]) -> New ordered dictionary with keys from S
                and values equal to v (which defaults to None).

                '''
                d = cls()
                for key in iterable:
                    d[key] = value
                return d

            def __eq__(self, other):
                '''od.__eq__(y) <==> od==y.  Comparison to another OD is order-sensitive
                while comparison to a regular mapping is order-insensitive.

                '''
                if isinstance(other, OrderedDict):
                    return len(self) == len(other) and self.items() == other.items()
                return dict.__eq__(self, other)

            def __ne__(self, other):
                return not self == other

#            # -- the following methods are only used in Python 2.7 --
#
#            def viewkeys(self):
#                "od.viewkeys() -> a set-like object providing a view on od's keys"
#                return KeysView(self)
#
#            def viewvalues(self):
#                "od.viewvalues() -> an object providing a view on od's values"
#                return ValuesView(self)
#
#            def viewitems(self):
#                "od.viewitems() -> a set-like object providing a view on od's items"
#                return ItemsView(self)
#        ## end of http://code.activestate.com/recipes/576693/ }}}

class DefaultOrderedDict(OrderedDict):
    'Dictionary that remembers insertion order and '
    def __init__(self, default_factory=None, *a, **kw):
        if (default_factory is not None and
            not isinstance(default_factory, Callable)):
            raise TypeError('first argument must be callable')
        super(DefaultOrderedDict, self).__init__(*a, **kw)
        self.default_factory = default_factory

    def __getitem__(self, key):
        try:
            return OrderedDict.__getitem__(self, key)
        except KeyError:
            return self.__missing__(key)

    def __missing__(self, key):
        if self.default_factory is None:
            raise KeyError(key)
        self[key] = value = self.default_factory()
        return value

    def __reduce__(self):
        if self.default_factory is None:
            args = tuple()
        else:
            args = self.default_factory,
        return type(self), args, None, None, self.items()

    def copy(self):
        return self.__copy__()

    def __copy__(self):
        return type(self)(self.default_factory, self)

    def __deepcopy__(self):
        import copy
        return type(self)(self.default_factory,
                          copy.deepcopy(self.items()))

    def __repr__(self, _repr_running={}):  # pylint: disable=W0102
        return 'DefaultOrderedDict({0}, {1})'.format(self.default_factory,
                                                     super(DefaultOrderedDict,
                                                           self).__repr__())
<commit_msg>import Callable from collections.abc for python3.7<commit_after># -*- coding: utf-8 -*-
'''
    :codeauthor: Pedro Algarvio (pedro@algarvio.me)

    salt.utils.odict
    ~~~~~~~~~~~~~~~~

    This is a compatibility/"importability" layer for an ordered dictionary.
    Tries to import from the standard library if python >= 2.7, then from the
    ``ordereddict`` package available from PyPi, and, as a last resort,
    provides an ``OrderedDict`` implementation based on::

        http://code.activestate.com/recipes/576669/

    It also implements a DefaultOrderedDict Class that serves  as a
    combination of ``OrderedDict`` and ``defaultdict``
    It's source was submitted here::

        http://stackoverflow.com/questions/6190331/
'''

# Import python libs
from __future__ import absolute_import, unicode_literals, print_function
try:
    from collections.abc import Callable
except ImportError:
    from collections import Callable

# Import 3rd-party libs
from salt.ext import six

try:
    # pylint: disable=E0611,minimum-python-version
    import collections

    class OrderedDict(collections.OrderedDict):
        __hash__ = None
except (ImportError, AttributeError):
    try:
        import ordereddict

        class OrderedDict(ordereddict.OrderedDict):  # pylint: disable=W0232
            __hash_ = None
    except ImportError:
        # {{{ http://code.activestate.com/recipes/576693/ (r9)
        # Backport of OrderedDict() class that runs on Python 2.4, 2.5, 2.6, 2.7 and pypy.
        # Passes Python2.7's test suite and incorporates all the latest updates.

        # pylint: disable=import-error,no-name-in-module
        try:
            from salt.ext.six.moves._thread import get_ident as _get_ident
        except ImportError:
            from salt.ext.six.moves._dummy_thread import get_ident as _get_ident
        # pylint: enable=import-error,no-name-in-module

#        try:
#            from _abcoll import KeysView, ValuesView, ItemsView
#        except ImportError:
#            pass

        class OrderedDict(dict):
            'Dictionary that remembers insertion order'
            # An inherited dict maps keys to values.
            # The inherited dict provides __getitem__, __len__, __contains__, and get.
            # The remaining methods are order-aware.
            # Big-O running times for all methods are the same as for regular dictionaries.

            # The internal self.__map dictionary maps keys to links in a doubly linked list.
            # The circular doubly linked list starts and ends with a sentinel element.
            # The sentinel element never gets deleted (this simplifies the algorithm).
            # Each link is stored as a list of length three:  [PREV, NEXT, KEY].
            __hash_ = None

            def __init__(self, *args, **kwds):  # pylint: disable=E1003
                '''Initialize an ordered dictionary.  Signature is the same as for
                regular dictionaries, but keyword arguments are not recommended
                because their insertion order is arbitrary.

                '''
                super(OrderedDict, self).__init__()  # pylint: disable=E1003
                if len(args) > 1:
                    raise TypeError(
                        'expected at most 1 arguments, got {0}'.format(len(args))
                    )
                try:
                    self.__root
                except AttributeError:
                    self.__root = root = []                     # sentinel node
                    root[:] = [root, root, None]
                    self.__map = {}
                self.__update(*args, **kwds)

            def __setitem__(self, key, value, dict_setitem=dict.__setitem__):
                'od.__setitem__(i, y) <==> od[i]=y'
                # Setting a new item creates a new link which goes at the end of the linked
                # list, and the inherited dictionary is updated with the new key/value pair.
                if key not in self:
                    root = self.__root
                    last = root[0]
                    last[1] = root[0] = self.__map[key] = [last, root, key]
                dict_setitem(self, key, value)

            def __delitem__(self, key, dict_delitem=dict.__delitem__):
                'od.__delitem__(y) <==> del od[y]'
                # Deleting an existing item uses self.__map to find the link which is
                # then removed by updating the links in the predecessor and successor nodes.
                dict_delitem(self, key)
                link_prev, link_next, key = self.__map.pop(key)
                link_prev[1] = link_next
                link_next[0] = link_prev

            def __iter__(self):
                'od.__iter__() <==> iter(od)'
                root = self.__root
                curr = root[1]
                while curr is not root:
                    yield curr[2]
                    curr = curr[1]

            def __reversed__(self):
                'od.__reversed__() <==> reversed(od)'
                root = self.__root
                curr = root[0]
                while curr is not root:
                    yield curr[2]
                    curr = curr[0]

            def clear(self):
                'od.clear() -> None.  Remove all items from od.'
                try:
                    for node in six.itervalues(self.__map):
                        del node[:]
                    root = self.__root
                    root[:] = [root, root, None]
                    self.__map.clear()
                except AttributeError:
                    pass
                dict.clear(self)

            def popitem(self, last=True):
                '''od.popitem() -> (k, v), return and remove a (key, value) pair.
                Pairs are returned in LIFO order if last is true or FIFO order if false.

                '''
                if not self:
                    raise KeyError('dictionary is empty')
                root = self.__root
                if last:
                    link = root[0]
                    link_prev = link[0]
                    link_prev[1] = root
                    root[0] = link_prev
                else:
                    link = root[1]
                    link_next = link[1]
                    root[1] = link_next
                    link_next[0] = root
                key = link[2]
                del self.__map[key]
                value = dict.pop(self, key)
                return key, value

            # -- the following methods do not depend on the internal structure --

            def keys(self):
                'od.keys() -> list of keys in od'
                return list(self)

            def values(self):
                'od.values() -> list of values in od'
                return [self[key] for key in self]

            def items(self):
                'od.items() -> list of (key, value) pairs in od'
                return [(key, self[key]) for key in self]

            def iterkeys(self):
                'od.iterkeys() -> an iterator over the keys in od'
                return iter(self)

            def itervalues(self):
                'od.itervalues -> an iterator over the values in od'
                for k in self:
                    yield self[k]

            def iteritems(self):
                'od.iteritems -> an iterator over the (key, value) items in od'
                for k in self:
                    yield (k, self[k])

            def update(*args, **kwds):  # pylint: disable=E0211
                '''od.update(E, **F) -> None.  Update od from dict/iterable E and F.

                If E is a dict instance, does:           for k in E: od[k] = E[k]
                If E has a .keys() method, does:         for k in E.keys(): od[k] = E[k]
                Or if E is an iterable of items, does:   for k, v in E: od[k] = v
                In either case, this is followed by:     for k, v in F.items(): od[k] = v

                '''
                if len(args) > 2:
                    raise TypeError(
                        'update() takes at most 2 positional '
                        'arguments ({0} given)'.format(len(args))
                    )
                elif not args:
                    raise TypeError('update() takes at least 1 argument (0 given)')
                self = args[0]
                # Make progressively weaker assumptions about "other"
                other = ()
                if len(args) == 2:
                    other = args[1]
                if isinstance(other, dict):
                    for key in other:
                        self[key] = other[key]
                elif hasattr(other, 'keys'):
                    for key in other:
                        self[key] = other[key]
                else:
                    for key, value in other:
                        self[key] = value
                for key, value in six.iteritems(kwds):
                    self[key] = value

            __update = update  # let subclasses override update without breaking __init__

            __marker = object()

            def pop(self, key, default=__marker):
                '''od.pop(k[,d]) -> v, remove specified key and return the corresponding value.
                If key is not found, d is returned if given, otherwise KeyError is raised.

                '''
                if key in self:
                    result = self[key]
                    del self[key]
                    return result
                if default is self.__marker:
                    raise KeyError(key)
                return default

            def setdefault(self, key, default=None):
                'od.setdefault(k[,d]) -> od.get(k,d), also set od[k]=d if k not in od'
                if key in self:
                    return self[key]
                self[key] = default
                return default

            def __repr__(self, _repr_running={}):  # pylint: disable=W0102
                'od.__repr__() <==> repr(od)'
                call_key = id(self), _get_ident()
                if call_key in _repr_running:
                    return '...'
                _repr_running[call_key] = 1
                try:
                    if not self:
                        return '{0}()'.format(self.__class__.__name__)
                    return '{0}(\'{1}\')'.format(self.__class__.__name__, list(self.items()))
                finally:
                    del _repr_running[call_key]

            def __reduce__(self):
                'Return state information for pickling'
                items = [[k, self[k]] for k in self]
                inst_dict = vars(self).copy()
                for k in vars(OrderedDict()):
                    inst_dict.pop(k, None)
                if inst_dict:
                    return (self.__class__, (items,), inst_dict)
                return self.__class__, (items,)

            def copy(self):
                'od.copy() -> a shallow copy of od'
                return self.__class__(self)

            @classmethod
            def fromkeys(cls, iterable, value=None):
                '''OD.fromkeys(S[, v]) -> New ordered dictionary with keys from S
                and values equal to v (which defaults to None).

                '''
                d = cls()
                for key in iterable:
                    d[key] = value
                return d

            def __eq__(self, other):
                '''od.__eq__(y) <==> od==y.  Comparison to another OD is order-sensitive
                while comparison to a regular mapping is order-insensitive.

                '''
                if isinstance(other, OrderedDict):
                    return len(self) == len(other) and self.items() == other.items()
                return dict.__eq__(self, other)

            def __ne__(self, other):
                return not self == other

#            # -- the following methods are only used in Python 2.7 --
#
#            def viewkeys(self):
#                "od.viewkeys() -> a set-like object providing a view on od's keys"
#                return KeysView(self)
#
#            def viewvalues(self):
#                "od.viewvalues() -> an object providing a view on od's values"
#                return ValuesView(self)
#
#            def viewitems(self):
#                "od.viewitems() -> a set-like object providing a view on od's items"
#                return ItemsView(self)
#        ## end of http://code.activestate.com/recipes/576693/ }}}

class DefaultOrderedDict(OrderedDict):
    'Dictionary that remembers insertion order and '
    def __init__(self, default_factory=None, *a, **kw):
        if (default_factory is not None and
            not isinstance(default_factory, Callable)):
            raise TypeError('first argument must be callable')
        super(DefaultOrderedDict, self).__init__(*a, **kw)
        self.default_factory = default_factory

    def __getitem__(self, key):
        try:
            return OrderedDict.__getitem__(self, key)
        except KeyError:
            return self.__missing__(key)

    def __missing__(self, key):
        if self.default_factory is None:
            raise KeyError(key)
        self[key] = value = self.default_factory()
        return value

    def __reduce__(self):
        if self.default_factory is None:
            args = tuple()
        else:
            args = self.default_factory,
        return type(self), args, None, None, self.items()

    def copy(self):
        return self.__copy__()

    def __copy__(self):
        return type(self)(self.default_factory, self)

    def __deepcopy__(self):
        import copy
        return type(self)(self.default_factory,
                          copy.deepcopy(self.items()))

    def __repr__(self, _repr_running={}):  # pylint: disable=W0102
        return 'DefaultOrderedDict({0}, {1})'.format(self.default_factory,
                                                     super(DefaultOrderedDict,
                                                           self).__repr__())
<|endoftext|>


<commit_before>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import socket

# save original socket method for restoration
# These are global so that re-calling the turn_off_internet function doesn't
# overwrite them again
socket_original = socket.socket
socket_create_connection = socket.create_connection
socket_bind = socket.socket.bind
socket_connect = socket.socket.connect

# ::1 is apparently another valid name for localhost?
# it is returned by getaddrinfo when that function is given localhost

def check_internet_off(original_function):
    def new_function(*args, **kwargs):
        if isinstance(args[0], socket.socket):
            host = args[1][0]
        else:
            host = args[0][0]
        if '127.0.0.1' in host or 'localhost' in host or '::1' in host:
            return original_function(*args, **kwargs)
        else:
            raise IOError("An attempt was made to connect to the internet "
                          "by a test that was not marked `remote_data`.")
    return new_function

def turn_off_internet(verbose=False):
    """
    Disable internet access via python by preventing connections from being
    created using the socket module.  Presumably this could be worked around by
    using some other means of accessing the internet, but all default python
    modules (urllib, requests, etc.) use socket [citation needed].
    """
    __tracebackhide__ = True
    if verbose:
        print("Internet access disabled")

    # ::1 is apparently another valid name for localhost?
    # it is returned by getaddrinfo when that function is given localhost

    def blocker_creator(return_fn, isclass=False):
        __tracebackhide__ = True
        def blocker(*args, **kwargs):
            if isclass:
                address = args[1]
            else:
                address = args[0]
            host,port = address[:2]
            if not ('127.0.0.1' in host or
                    'localhost' in host or
                    '::1' in host):
                raise IOError("An attempt was made to connect to the internet")
            else:
                return return_fn(*args, **kwargs)
        return blocker

    setattr(socket.socket, 'bind', blocker_creator(socket_bind, isclass=True))
    setattr(socket.socket, 'connect', blocker_creator(socket_connect, isclass=True))
    setattr(socket, 'create_connection', blocker_creator(socket_create_connect, isclass=False))

    return socket

def turn_on_internet(verbose=False):
    """
    Restore internet access.  Not used, but kept in case it is needed.
    """
    if verbose:
        print("Internet access enabled")
    socket.create_connection = socket_create_connection
    socket.socket.bind = socket_bind
    socket.socket.connect = socket_connect
    return socket
<commit_msg>fix to a bad rebase: the long, ugly version was NOT meant to be included<commit_after>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import socket

# save original socket method for restoration
# These are global so that re-calling the turn_off_internet function doesn't
# overwrite them again
socket_original = socket.socket
socket_create_connection = socket.create_connection
socket_bind = socket.socket.bind
socket_connect = socket.socket.connect

# ::1 is apparently another valid name for localhost?
# it is returned by getaddrinfo when that function is given localhost

def check_internet_off(original_function):
    def new_function(*args, **kwargs):
        if isinstance(args[0], socket.socket):
            host = args[1][0]
            valid_hosts = ('localhost', '127.0.0.1', '::1')
        else:
            host = args[0][0]
            valid_hosts = ('localhost', '127.0.0.1')
        if any([h in host for h in valid_hosts]):
            return original_function(*args, **kwargs)
        else:
            raise IOError("An attempt was made to connect to the internet "
                          "by a test that was not marked `remote_data`.")
    return new_function

def turn_off_internet(verbose=False):
    """
    Disable internet access via python by preventing connections from being
    created using the socket module.  Presumably this could be worked around by
    using some other means of accessing the internet, but all default python
    modules (urllib, requests, etc.) use socket [citation needed].
    """
    __tracebackhide__ = True
    if verbose:
        print("Internet access disabled")

    socket.create_connection = check_internet_off(socket_create_connection)
    socket.socket.bind = check_internet_off(socket_bind)
    socket.socket.connect = check_internet_off(socket_connect)

    return socket

def turn_on_internet(verbose=False):
    """
    Restore internet access.  Not used, but kept in case it is needed.
    """
    if verbose:
        print("Internet access enabled")
    socket.create_connection = socket_create_connection
    socket.socket.bind = socket_bind
    socket.socket.connect = socket_connect
    return socket
<|endoftext|>