<commit_before>import sys
import argparse
import os

from . import run
from . import get
from . import put
from . import ping
from . import ssh_copy_id

def get_version(include_path=False):
    import pkg_resources
    path = os.path.split(__file__)[0]
    try:
        pkg = pkg_resources.get_distribution('curlbomb')
    except pkg_resources.DistributionNotFound:
        return 'DEV' + ( " - {}".format(path) if include_path else "")
    if __file__.startswith(pkg.location):
        return pkg.version + ( " - {}".format(path) if include_path else "")
    return 'DEV' + ( " - {}".format(path) if include_path else "")

def add_inheritible_args(parser, subcommand=None):
    """Add the root curlbomb arguments to a (sub)parser

    specify subcommand=None for the root parser, substitute the
    subcommand name for subparsers.

    This allows core arguments to be specified either before or after
    the subparser name:

        curlbomb --ssh example.com:8080 run

    is equivalent to:

        curlbomb run --ssh example.com:8080

    Argparser default values, when combined with subparsers, is very
    tricky. You must follow these rules:

      * Arguments defined in this function MUST NOT have a default
        value (other than None). Values defined as store_true, must
        specify their default as None. Defaults are specified in
        settings:get_settings(). Without this rule, default values of
        the subparser will always take precedence over the parent
        parser value, *even if it's specified on the command line!*

        eg: curlbomb -n 5 run script.sh (If the run subcommand
        specifies n defaults to 1, this command is actually mistakenly
        run as -n 1)

      * Arguments must have a dest name prefixed with their subcommand
        name followed by a '_' when inherited into a subparser
        (eg. "run_ssh" for the run version of --ssh). This further
        decouples the subparser version from the root parser, and
        enables you to tell which one was used (by inspecting each
        value and seeing if it's None)

        This requires you to always check BOTH the main value as well
        as the sub_value (args.ssh and args.run_ssh). This logic is
        encapsulated in get_arg_value()

    """
    def dest(name):
        return ("" if not subcommand else subcommand + "_") + name
    server = parser.add_argument_group("These args modify the server")
    server.add_argument('-n', '--num-gets', metavar="N",
                        help="Number of times to serve resource (default:1)",
                        type=int, dest=dest("num_gets"))
    server.add_argument('-p', '--port',  metavar="PORT", help="TCP port number to use "
                        "(default:random available)",
                        dest=dest("port"))
    server.add_argument('--ssh', metavar="SSH_FORWARD", dest=dest("ssh"),
                        help="Forward curlbomb through another host via SSH - "
                        "[user@]host[:ssh_port][:http_port]")
    server.add_argument('--ssl', metavar="CERTIFICATE",
                        help="Use SSL with the given certificate file "
                        "(optionally PGP encrypted.) CERTIFICATE may be specified as a "
                        "single - to generate a new self-signed certificate and to turn "
                        "on --pin", default=None, dest=dest("ssl"))
    server.add_argument("-e", "--encrypt", action="store_true",
                        help="Encrypt resource with GPG before serving. Will use a randomly generated "
                        "symmetric passphrase unless --encrypt-to or --passphrase is specified.",
                        dest=dest("encrypt"), default=None)
    server.add_argument("--encrypt-to", action="append", metavar="GPG_ID", help="Encrypt with the "
                        "public key specified instead of a passphrase. Can be specified multiple times.",
                        dest=dest("encrypt_to"))
    server.add_argument("--passphrase", action="store_true",
                        help="Ask for a symmetric passphrase to encrypt with instead of a random one.",
                        dest=dest("passphrase"), default=None)
    server.add_argument('--mime-type', metavar="MIME/TYPE", help="The content type to serve",
                        dest=dest("mime_type"))
    server.add_argument('--disable-knock', action="store_true",
                        help="Don't require authentication (no X-knock header)",
                        dest=dest("disable_knock"), default=None)
    server.add_argument('--knock', metavar="KNOCK", help="Use a specific knock rather than random",
                        dest=dest("knock"))

    client = parser.add_argument_group("These args modify the client command")
    client.add_argument('-d','--domain', metavar="host[:port]",
                        help="Provide the domain and port to display "
                        "in the constructed URL. (example.com:8080)",
                        dest=dest("domain"))
    client.add_argument('-w', '--wget',
                        help="Output wget command rather than curl",
                        action="store_true", dest=dest("wget"), default=None)
    client.add_argument('--survey', help="Just a survey mission, no bomb run "
                        "(just get the script, don't run it)", action="store_true",
                        dest=dest("survey"), default=None)
    client.add_argument('--unwrapped',
                        help="Get the unwrapped version of the curlbomb "
                        "(1 less server request, but longer command)", action="store_true",
                        dest=dest("unwrapped"), default=None)
    client.add_argument('--pin', help="Pin the SSL certificate hash into the client "
                        "command to force curl to use our certificate"
                        " (requires --ssl)", action="store_true", dest=dest("pin"),
                        default=None)

    cli = parser.add_argument_group("These args modify CLI interaction")
    cli.add_argument('-l','--log-posts', dest=dest("log_post_backs"), action="store_true",
                        help="Log client stdout to server stdout", default=None)
    cli.add_argument('-q', '--quiet', action="store_true",
                        help="Be more quiet. Don't print the curlbomb command",
                        dest=dest("quiet"), default=None)
    cli.add_argument('-v', '--verbose', action="store_true",
                        help="Be more verbose. Enables --log-posts and print INFO logging",
                        dest=dest("verbose"), default=None)
    cli.add_argument('--log', metavar="LOG_FILE", dest=dest("log_file"),
                        help="log messages to LOG_FILE instead of stdout")
    cli.add_argument('--client-logging', dest=dest("client_logging"),
                        help="Enable client execution log (curlbomb.log on client)",
                        action="store_true", default=None)
    cli.add_argument('--client-quiet', dest=dest("client_quiet"),
                        help="Quiet the output on the client",
                        action="store_true", default=None)
    cli.add_argument('--pipe', help="Pipe to shell command rather than doing process substitution. "
                        "This is necessary for most interactive scripts.",
                        action="store_true", dest=dest("pipe"), default=None)
    cli.add_argument('--debug', action="store_true", default=None,
                        # Be really verbose, turn on debug logging
                        help=argparse.SUPPRESS, dest=dest("debug"))
    cli.add_argument('--version', action="version", version=get_version(True))

def argparser(formatter_class=argparse.HelpFormatter):
    parser = argparse.ArgumentParser(
        description='curlbomb is an HTTP server for serving one-time-use shell scripts',
        formatter_class=formatter_class)
    parser.set_defaults(subcommand=None)
    subparsers = parser.add_subparsers()

    add_inheritible_args(parser)

    # Old disabled method for disabling postback logging:
    # Now this just outputs an error if the user uses it:
    parser.add_argument('-1', '--disable-postback',
                        action="store_true", help=argparse.SUPPRESS, default=None)

    run.add_parser(subparsers)
    put.add_parser(subparsers)
    get.add_parser(subparsers)
    ping.add_parser(subparsers)
    ssh_copy_id.add_parser(subparsers)

    return parser

<commit_msg>Moves --pipe to client args group<commit_after>import sys
import argparse
import os

from . import run
from . import get
from . import put
from . import ping
from . import ssh_copy_id

def get_version(include_path=False):
    import pkg_resources
    path = os.path.split(__file__)[0]
    try:
        pkg = pkg_resources.get_distribution('curlbomb')
    except pkg_resources.DistributionNotFound:
        return 'DEV' + ( " - {}".format(path) if include_path else "")
    if __file__.startswith(pkg.location):
        return pkg.version + ( " - {}".format(path) if include_path else "")
    return 'DEV' + ( " - {}".format(path) if include_path else "")

def add_inheritible_args(parser, subcommand=None):
    """Add the root curlbomb arguments to a (sub)parser

    specify subcommand=None for the root parser, substitute the
    subcommand name for subparsers.

    This allows core arguments to be specified either before or after
    the subparser name:

        curlbomb --ssh example.com:8080 run

    is equivalent to:

        curlbomb run --ssh example.com:8080

    Argparser default values, when combined with subparsers, is very
    tricky. You must follow these rules:

      * Arguments defined in this function MUST NOT have a default
        value (other than None). Values defined as store_true, must
        specify their default as None. Defaults are specified in
        settings:get_settings(). Without this rule, default values of
        the subparser will always take precedence over the parent
        parser value, *even if it's specified on the command line!*

        eg: curlbomb -n 5 run script.sh (If the run subcommand
        specifies n defaults to 1, this command is actually mistakenly
        run as -n 1)

      * Arguments must have a dest name prefixed with their subcommand
        name followed by a '_' when inherited into a subparser
        (eg. "run_ssh" for the run version of --ssh). This further
        decouples the subparser version from the root parser, and
        enables you to tell which one was used (by inspecting each
        value and seeing if it's None)

        This requires you to always check BOTH the main value as well
        as the sub_value (args.ssh and args.run_ssh). This logic is
        encapsulated in get_arg_value()

    """
    def dest(name):
        return ("" if not subcommand else subcommand + "_") + name
    server = parser.add_argument_group("These args modify the server")
    server.add_argument('-n', '--num-gets', metavar="N",
                        help="Number of times to serve resource (default:1)",
                        type=int, dest=dest("num_gets"))
    server.add_argument('-p', '--port',  metavar="PORT", help="TCP port number to use "
                        "(default:random available)",
                        dest=dest("port"))
    server.add_argument('--ssh', metavar="SSH_FORWARD", dest=dest("ssh"),
                        help="Forward curlbomb through another host via SSH - "
                        "[user@]host[:ssh_port][:http_port]")
    server.add_argument('--ssl', metavar="CERTIFICATE",
                        help="Use SSL with the given certificate file "
                        "(optionally PGP encrypted.) CERTIFICATE may be specified as a "
                        "single - to generate a new self-signed certificate and to turn "
                        "on --pin", default=None, dest=dest("ssl"))
    server.add_argument("-e", "--encrypt", action="store_true",
                        help="Encrypt resource with GPG before serving. Will use a randomly generated "
                        "symmetric passphrase unless --encrypt-to or --passphrase is specified.",
                        dest=dest("encrypt"), default=None)
    server.add_argument("--encrypt-to", action="append", metavar="GPG_ID", help="Encrypt with the "
                        "public key specified instead of a passphrase. Can be specified multiple times.",
                        dest=dest("encrypt_to"))
    server.add_argument("--passphrase", action="store_true",
                        help="Ask for a symmetric passphrase to encrypt with instead of a random one.",
                        dest=dest("passphrase"), default=None)
    server.add_argument('--mime-type', metavar="MIME/TYPE", help="The content type to serve",
                        dest=dest("mime_type"))
    server.add_argument('--disable-knock', action="store_true",
                        help="Don't require authentication (no X-knock header)",
                        dest=dest("disable_knock"), default=None)
    server.add_argument('--knock', metavar="KNOCK", help="Use a specific knock rather than random",
                        dest=dest("knock"))

    client = parser.add_argument_group("These args modify the client command")
    client.add_argument('-d','--domain', metavar="host[:port]",
                        help="Provide the domain and port to display "
                        "in the constructed URL. (example.com:8080)",
                        dest=dest("domain"))
    client.add_argument('-w', '--wget',
                        help="Output wget command rather than curl",
                        action="store_true", dest=dest("wget"), default=None)
    client.add_argument('--survey', help="Just a survey mission, no bomb run "
                        "(just get the script, don't run it)", action="store_true",
                        dest=dest("survey"), default=None)
    client.add_argument('--unwrapped',
                        help="Get the unwrapped version of the curlbomb "
                        "(1 less server request, but longer command)", action="store_true",
                        dest=dest("unwrapped"), default=None)
    client.add_argument('--pin', help="Pin the SSL certificate hash into the client "
                        "command to force curl to use our certificate"
                        " (requires --ssl)", action="store_true", dest=dest("pin"),
                        default=None)
    client.add_argument('--pipe', help="Pipe to shell command rather than doing process substitution. "
                        "This is necessary for most interactive scripts.",
                        action="store_true", dest=dest("pipe"), default=None)

    cli = parser.add_argument_group("These args modify CLI interaction")
    cli.add_argument('-l','--log-posts', dest=dest("log_post_backs"), action="store_true",
                        help="Log client stdout to server stdout", default=None)
    cli.add_argument('-q', '--quiet', action="store_true",
                        help="Be more quiet. Don't print the curlbomb command",
                        dest=dest("quiet"), default=None)
    cli.add_argument('-v', '--verbose', action="store_true",
                        help="Be more verbose. Enables --log-posts and print INFO logging",
                        dest=dest("verbose"), default=None)
    cli.add_argument('--log', metavar="LOG_FILE", dest=dest("log_file"),
                        help="log messages to LOG_FILE instead of stdout")
    cli.add_argument('--client-logging', dest=dest("client_logging"),
                        help="Enable client execution log (curlbomb.log on client)",
                        action="store_true", default=None)
    cli.add_argument('--client-quiet', dest=dest("client_quiet"),
                        help="Quiet the output on the client",
                        action="store_true", default=None)
    cli.add_argument('--debug', action="store_true", default=None,
                        # Be really verbose, turn on debug logging
                        help=argparse.SUPPRESS, dest=dest("debug"))
    cli.add_argument('--version', action="version", version=get_version(True))

def argparser(formatter_class=argparse.HelpFormatter):
    parser = argparse.ArgumentParser(
        description='curlbomb is an HTTP server for serving one-time-use shell scripts',
        formatter_class=formatter_class)
    parser.set_defaults(subcommand=None)
    subparsers = parser.add_subparsers()

    add_inheritible_args(parser)

    # Old disabled method for disabling postback logging:
    # Now this just outputs an error if the user uses it:
    parser.add_argument('-1', '--disable-postback',
                        action="store_true", help=argparse.SUPPRESS, default=None)

    run.add_parser(subparsers)
    put.add_parser(subparsers)
    get.add_parser(subparsers)
    ping.add_parser(subparsers)
    ssh_copy_id.add_parser(subparsers)

    return parser

<|endoftext|>


<commit_before>"""
    Tests for the h5py.Datatype class.
"""

from __future__ import absolute_import

from itertools import count
import numpy as np
import h5py

from ..common import ut, TestCase

class TestVlen(TestCase):

    """
        Check that storage of vlen strings is carried out correctly.
    """
    def assertVlenArrayEqual(self, dset, arr, message=None, precision=None):
        self.assert_(
            dset.shape == arr.shape,
            "Shape mismatch (%s vs %s)%s" % (dset.shape, arr.shape, message)
            )
        for (i, d, a) in zip(count(), dset, arr):
            self.assertArrayEqual(d, a, message, precision)

    def test_compound(self):

        fields = []
        fields.append(('field_1', h5py.special_dtype(vlen=str)))
        fields.append(('field_2', np.int32))
        dt = np.dtype(fields)
        self.f['mytype'] = np.dtype(dt)
        dt_out = self.f['mytype'].dtype.fields['field_1'][0]
        self.assertEqual(h5py.check_dtype(vlen=dt_out), str)

    def test_compound_bool(self):
        vidt = h5py.special_dtype(vlen=np.uint8)
        eidt = h5py.special_dtype(enum=(np.uint8, {'OFF': 0, 'ON': 1}))
        f = self.f

        dt_vb = np.dtype([
            ('foo', vidt),
            ('logical', np.bool)])
        vb = f.create_dataset('dt_vb', shape=(4,), dtype=dt_vb)
        data = np.array([(np.array([1,2,3], dtype=np.uint8), True),
                         (np.array([1    ], dtype=np.uint8), False),
                         (np.array([1,5  ], dtype=np.uint8), True),
                         (np.array([],      dtype=np.uint8), False),],
                     dtype=dt_vb)
        vb[:] = data
        actual = f['dt_vb'][:]
        self.assertVlenArrayEqual(data['foo'], actual['foo'])
        self.assertArrayEqual(data['logical'], actual['logical'])

        dt_vv = np.dtype([
            ('foo', vidt),
            ('bar', vidt)])
        f.create_dataset('dt_vv', shape=(4,), dtype=dt_vv)

        dt_vve = np.dtype([
            ('foo', vidt),
            ('bar', vidt),
            ('switch', eidt)])
        vve = f.create_dataset('dt_vve', shape=(2,), dtype=dt_vve)
        data = np.array([(np.array([1,2,3]), np.array([1,2]),   1),
                         (np.array([]),      np.array([2,4,6]), 0),],
                         dtype=dt_vve)
        vve[:] = data

        dt_vvb = np.dtype([
            ('foo', vidt),
            ('bar', vidt),
            ('logical', np.bool)])
        vvb = f.create_dataset('dt_vvb', shape=(2,), dtype=dt_vvb)

    def test_vlen_enum(self):
        fname = self.mktemp()
        arr1 = [[1],[1,2]]
        dt1 = h5py.special_dtype(vlen=h5py.special_dtype(
            enum=('i', dict(foo=1, bar=2))))

        with h5py.File(fname,'w') as f:
            df1 = f.create_dataset('test', (len(arr1),), dtype=dt1)
            df1[:] = np.array(arr1)

        with h5py.File(fname,'r') as f:
            df2  = f['test']
            dt2  = df2.dtype
            arr2 = [e.tolist() for e in df2[:]]

        self.assertEqual(arr1, arr2)
        self.assertEqual(h5py.check_dtype(enum=h5py.check_dtype(vlen=dt1)),
                         h5py.check_dtype(enum=h5py.check_dtype(vlen=dt2)))

class TestOffsets(TestCase):
    """
        Check that compound members with aligned or manual offsets are handled
        correctly.
    """

    def test_compound_vlen(self):
        vidt = h5py.special_dtype(vlen=np.uint8)
        eidt = h5py.special_dtype(enum=(np.uint8, {'OFF': 0, 'ON': 1}))

        for np_align in (False, True):
            dt = np.dtype([
                ('foo', vidt),
                ('bar', vidt),
                ('switch', eidt)], align=np_align)

            for logical in (False, True):
                if np_align:
                    # Vlen types have different size in the numpy struct
                    self.assertRaises(TypeError, h5py.h5t.py_create, dt,
                            logical=logical)
                else:
                    ht = h5py.h5t.py_create(dt, logical=logical)

    def test_aligned_offsets(self):
        dt = np.dtype('i2,i8', align=True)
        ht = h5py.h5t.py_create(dt)
        self.assertEqual(dt.itemsize, ht.get_size())
        self.assertEqual(
            [dt.fields[i][1] for i in dt.names],
            [ht.get_member_offset(i) for i in range(ht.get_nmembers())]
        )

    def test_aligned_data(self):
        dt = np.dtype('i2,f8', align=True)
        data = np.empty(10, dtype=dt)

        data['f0'] = np.array(np.random.randint(-100, 100, size=data.size),
                dtype='i2')
        data['f1'] = np.random.rand(data.size)

        fname = self.mktemp()

        with h5py.File(fname, 'w') as f:
            f['data'] = data

        with h5py.File(fname, 'r') as f:
            self.assertArrayEqual(f['data'], data)

    def test_out_of_order_offsets(self):
        dt = np.dtype({
            'names' : ['f1', 'f2', 'f3'],
            'formats' : ['<f4', '<i4', '<f8'],
            'offsets' : [0, 16, 8]
        })
        data = np.empty(10, dtype=dt)
        data['f1'] = np.random.rand(data.size)
        data['f2'] = np.random.random_integers(-10, 10, data.size)
        data['f3'] = np.random.rand(data.size)*-1

        fname = self.mktemp()

        with h5py.File(fname, 'w') as fd:
            fd.create_dataset('data', data=data)

        with h5py.File(fname, 'r') as fd:
            self.assertArrayEqual(fd['data'], data)
<commit_msg>Add additional test<commit_after>"""
    Tests for the h5py.Datatype class.
"""

from __future__ import absolute_import

from itertools import count
import numpy as np
import h5py

from ..common import ut, TestCase

class TestVlen(TestCase):

    """
        Check that storage of vlen strings is carried out correctly.
    """
    def assertVlenArrayEqual(self, dset, arr, message=None, precision=None):
        self.assert_(
            dset.shape == arr.shape,
            "Shape mismatch (%s vs %s)%s" % (dset.shape, arr.shape, message)
            )
        for (i, d, a) in zip(count(), dset, arr):
            self.assertArrayEqual(d, a, message, precision)

    def test_compound(self):

        fields = []
        fields.append(('field_1', h5py.special_dtype(vlen=str)))
        fields.append(('field_2', np.int32))
        dt = np.dtype(fields)
        self.f['mytype'] = np.dtype(dt)
        dt_out = self.f['mytype'].dtype.fields['field_1'][0]
        self.assertEqual(h5py.check_dtype(vlen=dt_out), str)

    def test_compound_bool(self):
        vidt = h5py.special_dtype(vlen=np.uint8)
        def a(items):
            return np.array(items, dtype=np.uint8)

        eidt = h5py.special_dtype(enum=(np.uint8, {'OFF': 0, 'ON': 1}))
        f = self.f

        dt_vb = np.dtype([
            ('foo', vidt),
            ('logical', np.bool)])
        vb = f.create_dataset('dt_vb', shape=(4,), dtype=dt_vb)
        data = np.array([(a([1,2,3]), True),
                         (a([1    ]), False),
                         (a([1,5  ]), True),
                         (a([],    ), False),],
                     dtype=dt_vb)
        vb[:] = data
        actual = f['dt_vb'][:]
        self.assertVlenArrayEqual(data['foo'], actual['foo'])
        self.assertArrayEqual(data['logical'], actual['logical'])

        dt_vv = np.dtype([
            ('foo', vidt),
            ('bar', vidt)])
        f.create_dataset('dt_vv', shape=(4,), dtype=dt_vv)

        dt_vve = np.dtype([
            ('foo', vidt),
            ('bar', vidt),
            ('switch', eidt)])
        vve = f.create_dataset('dt_vve', shape=(2,), dtype=dt_vve)
        data = np.array([(a([1,2,3]), a([1,2]),   1),
                         (a([]),      a([2,4,6]), 0),],
                         dtype=dt_vve)
        vve[:] = data
        actual = vve[:]
        self.assertVlenArrayEqual(data['foo'], actual['foo'])
        self.assertVlenArrayEqual(data['bar'], actual['bar'])
        self.assertArrayEqual(data['switch'], actual['switch'])

        dt_vvb = np.dtype([
            ('foo', vidt),
            ('bar', vidt),
            ('logical', np.bool)])
        vvb = f.create_dataset('dt_vvb', shape=(2,), dtype=dt_vvb)

    def test_vlen_enum(self):
        fname = self.mktemp()
        arr1 = [[1],[1,2]]
        dt1 = h5py.special_dtype(vlen=h5py.special_dtype(
            enum=('i', dict(foo=1, bar=2))))

        with h5py.File(fname,'w') as f:
            df1 = f.create_dataset('test', (len(arr1),), dtype=dt1)
            df1[:] = np.array(arr1)

        with h5py.File(fname,'r') as f:
            df2  = f['test']
            dt2  = df2.dtype
            arr2 = [e.tolist() for e in df2[:]]

        self.assertEqual(arr1, arr2)
        self.assertEqual(h5py.check_dtype(enum=h5py.check_dtype(vlen=dt1)),
                         h5py.check_dtype(enum=h5py.check_dtype(vlen=dt2)))

class TestOffsets(TestCase):
    """
        Check that compound members with aligned or manual offsets are handled
        correctly.
    """

    def test_compound_vlen(self):
        vidt = h5py.special_dtype(vlen=np.uint8)
        eidt = h5py.special_dtype(enum=(np.uint8, {'OFF': 0, 'ON': 1}))

        for np_align in (False, True):
            dt = np.dtype([
                ('foo', vidt),
                ('bar', vidt),
                ('switch', eidt)], align=np_align)

            for logical in (False, True):
                if np_align:
                    # Vlen types have different size in the numpy struct
                    self.assertRaises(TypeError, h5py.h5t.py_create, dt,
                            logical=logical)
                else:
                    ht = h5py.h5t.py_create(dt, logical=logical)

    def test_aligned_offsets(self):
        dt = np.dtype('i2,i8', align=True)
        ht = h5py.h5t.py_create(dt)
        self.assertEqual(dt.itemsize, ht.get_size())
        self.assertEqual(
            [dt.fields[i][1] for i in dt.names],
            [ht.get_member_offset(i) for i in range(ht.get_nmembers())]
        )

    def test_aligned_data(self):
        dt = np.dtype('i2,f8', align=True)
        data = np.empty(10, dtype=dt)

        data['f0'] = np.array(np.random.randint(-100, 100, size=data.size),
                dtype='i2')
        data['f1'] = np.random.rand(data.size)

        fname = self.mktemp()

        with h5py.File(fname, 'w') as f:
            f['data'] = data

        with h5py.File(fname, 'r') as f:
            self.assertArrayEqual(f['data'], data)

    def test_out_of_order_offsets(self):
        dt = np.dtype({
            'names' : ['f1', 'f2', 'f3'],
            'formats' : ['<f4', '<i4', '<f8'],
            'offsets' : [0, 16, 8]
        })
        data = np.empty(10, dtype=dt)
        data['f1'] = np.random.rand(data.size)
        data['f2'] = np.random.random_integers(-10, 10, data.size)
        data['f3'] = np.random.rand(data.size)*-1

        fname = self.mktemp()

        with h5py.File(fname, 'w') as fd:
            fd.create_dataset('data', data=data)

        with h5py.File(fname, 'r') as fd:
            self.assertArrayEqual(fd['data'], data)
<|endoftext|>


<commit_before>import json

from common.web import requestsManager
from constants import exceptions
from objects import glob

class handler(requestsManager.asyncRequestHandler):
	def asyncGet(self):
		statusCode = 400
		data = {"message": "unknown error"}
		try:
			# Check arguments
			if "u" not in self.request.arguments and "id" not in self.request.arguments:
				raise exceptions.invalidArgumentsException()

			# Get online staus
			username = None
			userID = None
			if "u" in self.request.arguments:
				username = self.get_argument("u").lower().replace(" ", "_")
			else:
				try:
					userID = int(self.get_argument("id"))
				except:
					raise exceptions.invalidArgumentsException()

			if username is None and userID is None:
				data["result"] = False
			else:
				if username is not None:
					data["result"] = True if glob.tokens.getTokenFromUsername(username, safe=True) is not None else False
				else:
					data["result"] = True if glob.tokens.getTokenFromUserID(userID) is not None else False

			# Status code and message
			statusCode = 200
			data["message"] = "ok"
		except exceptions.invalidArgumentsException:
			statusCode = 400
			data["message"] = "missing required arguments"
		finally:
			# Add status code to data
			data["status"] = statusCode

			# Send response
			self.write(json.dumps(data))
			self.set_status(statusCode)
<commit_msg>.BANCHO. .FIX. Run safeUsername on /api/v1/isOnline<commit_after>import json

from common.ripple import userUtils
from common.web import requestsManager
from constants import exceptions
from objects import glob

class handler(requestsManager.asyncRequestHandler):
	def asyncGet(self):
		statusCode = 400
		data = {"message": "unknown error"}
		try:
			# Check arguments
			if "u" not in self.request.arguments and "id" not in self.request.arguments:
				raise exceptions.invalidArgumentsException()

			# Get online staus
			username = None
			userID = None
			if "u" in self.request.arguments:
				#username = self.get_argument("u").lower().replace(" ", "_")
				username = userUtils.safeUsername(self.get_argument("u"))
			else:
				try:
					userID = int(self.get_argument("id"))
				except:
					raise exceptions.invalidArgumentsException()

			if username is None and userID is None:
				data["result"] = False
			else:
				if username is not None:
					data["result"] = True if glob.tokens.getTokenFromUsername(username, safe=True) is not None else False
				else:
					data["result"] = True if glob.tokens.getTokenFromUserID(userID) is not None else False

			# Status code and message
			statusCode = 200
			data["message"] = "ok"
		except exceptions.invalidArgumentsException:
			statusCode = 400
			data["message"] = "missing required arguments"
		finally:
			# Add status code to data
			data["status"] = statusCode

			# Send response
			self.write(json.dumps(data))
			self.set_status(statusCode)
<|endoftext|>


<commit_before>import logging, pprint

from commands import command

logger = logging.getLogger(__name__)

def _initialise(bot): pass # prevents commands from being automatically added

@command.register(admin=True)
def tagset(bot, event, *args):
    if len(args) == 3:
        [type, id, tag] = args
        if bot.tags.add(type, id, tag):
            message = _("tagged
{}
 with
{}
".format(id, tag))
        else:
            message = _("
{}
 unchanged".format(id))
    else:
        message = _("supply type, id, tag")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagdel(bot, event, *args):
    if len(args) == 3:
        [type, id, tag] = args
        if bot.tags.remove(type, id, tag):
            message = _("removed
{}
 from
{}
".format(tag, id))
        else:
            message = _("
{}
 unchanged".format(id))
    else:
        message = _("supply type, id, tag")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagspurge(bot, event, *args):
    if len(args) == 2:
        [type, id] = args
        entries_removed = bot.tags.purge(type, id)
        message = _("entries removed:
{}
".format(entries_removed))
    else:
        message = _("supply type, id")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagindexdump(bot, event, *args):
    pp = pprint.PrettyPrinter(indent=2)
    pp.pprint(bot.tags.indices)

    chunks = []
    for relationship in bot.tags.indices:
        lines = [_("index:
{}
").format(relationship)]
        for key, list in bot.tags.indices[relationship].items():
            lines.append(_("key:
{}
").format(key))
            for item in list:
                lines.append("...
{}
".format(item))
        if len(lines) == 0:
            continue
        chunks.append("
".join(lines))

    if len(chunks) == 0:
        chunks = [_("no entries to list")]

    bot.send_message_parsed(event.conv_id, "

".join(chunks))

@command.register(admin=True)
def tagsuser(bot, event, *args):
    if len(args) == 1:
        conv_id = "*"
        chat_id = args[0]
    elif len(args) == 2:
        conv_id = args[1]
        chat_id = args[0]
    else:
        bot.send_message_parsed(event.conv_id, _("supply chat_id, optional conv_id"))
        return

    active_user_tags = bot.tags.useractive(chat_id, conv_id)
    if not active_user_tags:
        active_user_tags = [_("no tags returned")]

    bot.send_message_parsed(event.conv_id, "
{}
@
{}
: {}".format(
        chat_id, conv_id, ", ".join(active_user_tags)))

@command.register(admin=True)
def tagsuserlist(bot, event, *args):
    if len(args) == 1:
        conv_id = args[0]
        filter_tags = False
    elif len(args) > 1:
        conv_id = args[0]
        filter_tags = args[1:]
    else:
        bot.send_message_parsed(event.conv_id, _("supply conv_id, optional tag list"))
        return

    users_to_tags = bot.tags.userlist(conv_id, filter_tags)

    lines = []
    for chat_id, active_user_tags in users_to_tags.items():
        if not active_user_tags:
            active_user_tags = [_("no tags returned")]
        lines.append("
{}
:
{}
".format(chat_id, ", ".join(active_user_tags)))

    if len(lines) == 0:
        lines = [_("no users found")]

    bot.send_message_parsed(event.conv_id, "
".join(lines))
<commit_msg>tagscommand to help figure out which tags apply to a command updated/added docstrings<commit_after>import logging, pprint

from commands import command

logger = logging.getLogger(__name__)

def _initialise(bot): pass # prevents commands from being automatically added

@command.register(admin=True)
def tagset(bot, event, *args):
    """set a single tag. usage: tagset <"conv"|"user"|"convuser"> <id> <tag>"""
    if len(args) == 3:
        [type, id, tag] = args
        if bot.tags.add(type, id, tag):
            message = _("tagged
{}
 with
{}
".format(id, tag))
        else:
            message = _("
{}
 unchanged".format(id))
    else:
        message = _("supply type, id, tag")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagdel(bot, event, *args):
    """remove single tag. usage: tagdel <"conv"|"user"|"convuser"> <id> <tag>"""
    if len(args) == 3:
        [type, id, tag] = args
        if bot.tags.remove(type, id, tag):
            message = _("removed
{}
 from
{}
".format(tag, id))
        else:
            message = _("
{}
 unchanged".format(id))
    else:
        message = _("supply type, id, tag")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagspurge(bot, event, *args):
    """batch remove tags. usage: tagspurge <"user"|"conv"|"convuser"|"tag"|"usertag"|"convtag"> <id|"ALL">"""
    if len(args) == 2:
        [type, id] = args
        entries_removed = bot.tags.purge(type, id)
        message = _("entries removed:
{}
".format(entries_removed))
    else:
        message = _("supply type, id")
    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagscommand(bot, event, *args):
    """display of command tagging information, more complete than plugininfo"""
    if len(args) == 1:
        [command_name] = args

        if command_name not in command.commands:
            message = _("
COMMAND: {}
 does not exist".format(command_name))

        else:
            lines = []

            ALL_TAGS = set()

            plugin_defined = set()
            if command_name in command.command_tagsets:
                plugin_defined = command.command_tagsets[command_name]
                ALL_TAGS = ALL_TAGS | plugin_defined

            config_root = set()
            config_commands_tagged = bot.get_config_option('commands_tagged') or {}
            if command_name in config_commands_tagged and config_commands_tagged[command_name]:
                config_root = set([ frozenset(value if isinstance(value, list) else [value])
                    for value in config_commands_tagged[command_name] ])
                ALL_TAGS = ALL_TAGS | config_root

            config_conv = {}
            if bot.config.exists(["conversations"]):
                for convid in bot.config["conversations"]:
                    if bot.config.exists(["conversations", convid, "commands_tagged"]):
                        conv_tagged = bot.config.get_by_path(["conversations", convid, "commands_tagged"])
                        if command_name in conv_tagged and conv_tagged[command_name]:
                            config_conv[convid] = set([ frozenset(value if isinstance(value, list) else [value])
                                for value in conv_tagged[command_name] ])
                            ALL_TAGS = ALL_TAGS | config_conv[convid]

            dict_tags = {}
            for match in ALL_TAGS:
                text_match = ", ".join(sorted(match))

                if text_match not in dict_tags:
                    dict_tags[text_match] = []

                if match in plugin_defined:
                    dict_tags[text_match].append("plugin")
                if match in config_root:
                    dict_tags[text_match].append("config: root")
                for convid, tagsets in config_conv.items():
                    if match in tagsets:
                        dict_tags[text_match].append("config: {}".format(convid))

            for text_tags in sorted(dict_tags.keys()):
                lines.append("[ {} ]".format(text_tags))
                for source in dict_tags[text_tags]:
                    lines.append("... {}".format(source))

            if len(lines)==0:
                message = _("
COMMAND: {}
 has no tags".format(command_name))
            else:
                lines.insert(0, _("
COMMAND: {}
, match ANY:".format(command_name)))
                message = "
".join(lines)

    else:
        message = _("supply command name")

    bot.send_message_parsed(event.conv_id, message)

@command.register(admin=True)
def tagindexdump(bot, event, *args):
    """dump raw contents of tags indices"""
    pp = pprint.PrettyPrinter(indent=2)
    pp.pprint(bot.tags.indices)

    chunks = []
    for relationship in bot.tags.indices:
        lines = [_("index:
{}
").format(relationship)]
        for key, list in bot.tags.indices[relationship].items():
            lines.append(_("key:
{}
").format(key))
            for item in list:
                lines.append("...
{}
".format(item))
        if len(lines) == 0:
            continue
        chunks.append("
".join(lines))

    if len(chunks) == 0:
        chunks = [_("no entries to list")]

    bot.send_message_parsed(event.conv_id, "

".join(chunks))

@command.register(admin=True)
def tagsuser(bot, event, *args):
    """get tag assignments for a user in an (optional) conversation. usage: tagsuser <user id> [<conv id>]"""
    if len(args) == 1:
        conv_id = "*"
        chat_id = args[0]
    elif len(args) == 2:
        conv_id = args[1]
        chat_id = args[0]
    else:
        bot.send_message_parsed(event.conv_id, _("supply chat_id, optional conv_id"))
        return

    active_user_tags = bot.tags.useractive(chat_id, conv_id)
    if not active_user_tags:
        active_user_tags = [_("no tags returned")]

    bot.send_message_parsed(event.conv_id, "
{}
@
{}
: {}".format(
        chat_id, conv_id, ", ".join(active_user_tags)))

@command.register(admin=True)
def tagsuserlist(bot, event, *args):
    """get tag assignments for all users in a conversation, filtered by (optional) taglist. usage: tagsuserlist <conv id> [<tag name> [<tag name>] [...]]"""
    if len(args) == 1:
        conv_id = args[0]
        filter_tags = False
    elif len(args) > 1:
        conv_id = args[0]
        filter_tags = args[1:]
    else:
        bot.send_message_parsed(event.conv_id, _("supply conv_id, optional tag list"))
        return

    users_to_tags = bot.tags.userlist(conv_id, filter_tags)

    lines = []
    for chat_id, active_user_tags in users_to_tags.items():
        if not active_user_tags:
            active_user_tags = [_("no tags returned")]
        lines.append("
{}
:
{}
".format(chat_id, ", ".join(active_user_tags)))

    if len(lines) == 0:
        lines = [_("no users found")]

    bot.send_message_parsed(event.conv_id, "
".join(lines))
<|endoftext|>


<commit_before>from __future__ import division
import numpy as np
import cmath
from scipy import pi, exp, sin
import scipy
from numpy.random import uniform
import itertools
from collections import deque
import logging

from .IterativeMethods import iterateToRoot
from .CountRoots import prod, RootError

def subdivide(boxDeque, parentBox, parentBox_numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax):
	numberOfRoots = None
	for subBoxes in parentBox.subdivisions():
		try:
			numberOfRoots = [box.count_roots(f, df, integerTol, integrandUpperBound, divMax) for box in np.array(subBoxes)]
			if parentBox_numberOfRoots == sum(numberOfRoots):
				break
		except RootError:
			# If the number of zeros within either of the new contours is not an integer then it is
			# likely that the introduced line which subdivides 'parentBox' lies on a zero.
			# To avoid this we will try to place the subdividing line at a different point along
			# the division axis
			continue

	if numberOfRoots is None or parentBox_numberOfRoots != sum(numberOfRoots):
		# The list of subdivisions has been exhaused and still the number of enclosed zeros does not add up
		raise RuntimeError("""Unable to subdivide box:
			\t%s
			Consider increasing the integrandUpperBound to allow contours closer to roots to be integrated.""" % parentBox)

	boxDeque.extend([(box, numberOfRoots[i]) for i, box in enumerate(subBoxes) if numberOfRoots[i] != 0])

def addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter):
	# check that the root we have found is distinct from the ones we already have
	if not roots or np.all(abs(np.array(roots) - root) > newtonStepTol):
		# add the root to the list if it is within the original box
		if originalContour.contains(root):
			roots.append(root)
			multiplicities.append(find_multiplicity(root, f, df, rootErrTol))

		# check to see if there are any other roots implied by the given symmetry
		if guessRootSymmetry is not None:
			for x0 in guessRootSymmetry(root):
				root = iterateToRoot(x0, f, df, newtonStepTol, rootErrTol, newtonMaxIter)
				if root is not None:
					addRoot(root, roots, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

def find_multiplicity(root, f, df=None, rootErrTol=1e-12, dx=1e-8):
	# given a root of the function f find determine the multiplicity
	# of the root numerically
	i, f_root = 0, f(root)
	while abs(f_root) < rootErrTol:
		i += 1
		if df == None:
			f_root = scipy.misc.derivative(f, root, dx=dx, n=i, order=2*i+1)
		else:
			if i == 1:
				f_root = df(root)
			else:
				f_root = scipy.misc.derivative(df, root, dx=dx, n=i-1, order=2*(i-1)+1)
	return i

def findRootsGen(originalContour, f, df=None, guessRoot=[], guessRootSymmetry=None,
	newtonStepTol=1e-8, newtonMaxIter=20, rootErrTol=1e-12,
	absTol=1e-12, relTol=1e-12, divMax=20, integerTol=0.25, integrandUpperBound=1e3,
	M=5):
	"""
	A generator which at each step takes a contour and either finds
	all the zeros of f within it or subdivides it further.  Based
	on the algorithm in [KB]

	Parameters
	----------
	originalContour : subclass of Contour
		The contour C which bounds the region in which all the
		roots of f(z) are sought.
	f : function
		A function of a single complex variable, z, which is
		analytic within C and has no poles or roots on C.
	df : function, optional
		A function of a single complex variable which is the
		derivative of the function f(z). If df is not given
		then it will be approximated with a finite difference
		formula.
	guessRoot : list, optional
		A list of (root, multiplicity) tuples giving the already
		known roots of the function f.
	guessRootSymmetry : function, optional
		A function of a single complex variable, z, which
		returns a list of all points which are expected to be
		roots of f, given that z is a root of f.  It is assumed
		that the symmetric roots have the same multiplicity of
		the original root.
	newtonStepTol : float, optional
		The iterative method used to give a final value for each
		root will exit if the step size, dx, between sucessive
		iterations satisfied abs(dx) < newtonStepTol
	newtonMaxIter : int, optional
		The iterative method used to give a final value for each
		root will exit if the number of iterations exceeds newtonMaxIter.
	rootErrTol : float, optional
		The iterative method used to give a final value for each
		root will exit if, for a point approimating the root, x,
		abs(f(x)) < rootErrTol.
	absTol : float, optional
		The integration along a path will return a result if the
		difference between the last two iterations is less than abstol.
	relTol : float, optional
		The integration along a path will return a result if the
		difference between the last two iterations are less than
		relTol multiplied by the last iteration.
	divMax : int, optional
		The maximum number of divisions before the Romberg integration
		routine of a path exits.
	integerTol : float, optional
		If the result of a contour integration is expected to be
		an integer (ie. when computing the number of zeros within
		a contour) then the result of the last iteration of the
		intergration must be within integerTol of an integer.
		The parameters absTol and relTol are not used in this case.
	integrandUpperBound : float, optional
		The maximum allowed value of abs(df(z)/f(z)) along a contour.
		If abs(df(z)/f(z)) is found to exceed this value then there
		is likely to be a root close to the contour and the
		contour is discarded and a new subdivision created.
		If integrandUpperBound is too large then integrals may take
		a very long time to converge and it is generally be more
		efficient to allow the rootfinding procedure to instead
		choose another contour then spend time evaluting the
		integral along a contour very close to a root.
	M : int, optional
		If the number of roots (including multiplicites) within a
		contour is greater than M then the contour is subdivided
		further.  M must be greater than or equal to the largest
		multiplcity of any root.

	Yields
	------
	list
		Roots of f(z) within the contour C
	list
		Multiplicites of roots
	deque
		The contours which still contain roots
	int
		Remaining number of roots to be found within the contour

	References
	----------
	[KB] Peter Kravanja, Marc Van Barel, "Computing the Zeros of
	Anayltic Functions", Springer (2000)
	"""
	try:
		# total number of zeros, including multiplicities
		totNumberOfRoots = originalContour.count_roots(f, df, integerTol, integrandUpperBound, divMax)
	except RuntimeError:
		raise RuntimeError("""
			Integration along the initial contour has failed.  There is likely a root on or close to the initial contour
			Try either changing the initial contour, if possible, or increasing the integrandUpperBound to allow for
			a longer integration time.""")

	loggedIterativeWarning = False
	roots = []
	multiplicities = []

	# Add given roots
	# XXX: check these roots and multiplcities
	for root, multiplicity in guessRoot:
		addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

	print('Tot number of Roots', totNumberOfRoots)

	boxes = deque()
	boxes.append((originalContour,totNumberOfRoots))
	while boxes:
		box, numberOfRoots = boxes.pop()

		if numberOfRoots > M:
			subdivide(boxes, box, numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)

		else:
			# approximate the roots in this box
			approxRoots, approxRootMultiplicities = box.approximate_roots(f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)
			for approxRoot, multiplicity in list(zip(approxRoots, approxRootMultiplicities)):
				# if all the roots within the box have been located then just exit
				if numberOfRoots == sum([int(round(multiplicity.real)) for root, multiplicity in zip(roots, multiplicities) if box.contains(root)]):
					break

				if abs(f(approxRoot)) < rootErrTol:
					# the approximate root is good enough
					root = approxRoot
				else:
					# attempt to refine the root
					root = iterateToRoot(approxRoot, f, df, newtonStepTol, rootErrTol, newtonMaxIter)

				if root is not None:
					# if we found a root add it to the list of known roots
					addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

			# if we haven't found all the roots then subdivide further
			if numberOfRoots != sum([int(round(multiplicity.real)) for root, multiplicity in zip(roots, multiplicities) if box.contains(root)]):
				subdivide(boxes, box, numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)

		yield roots, multiplicities, boxes, totNumberOfRoots - len(roots)

	if totNumberOfRoots == 0:
		yield [], [], deque(), 0

def findRoots(originalContour, f, df=None, **kwargs):
	"""
	Return a list of all roots of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.
	"""
	rootFinder = findRootsGen(originalContour, f, df, **kwargs)
	for roots, multiplicities, boxes, numberOfRemainingRoots in rootFinder:
		pass
	return roots, multiplicities

def demo_findRoots(originalContour, f, df=None, automaticAnimation=False, returnAnim=False, **kwargs):
	"""
	An interactive demonstration of the processess used to find all the roots
	of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.

	If automaticAnimation is False (default) then press the SPACE key
	to step the animation forward.

	If automaticAnimation is True then the animation will play automatically
	until all the roots have been found.

	If returnAnim is true the animating object returned by matplotlib's animation.FuncAnimation
	will be returned, rather than the animation be shown.
	"""
	import matplotlib.pyplot as plt
	from matplotlib import animation
	rootFinder = findRootsGen(originalContour, f, df, **kwargs)

	originalContour.plot(linecolor='k', linestyle='--')

	fig = plt.gcf()
	ax = plt.gca()

	def update_frame(args):
		roots, multiplicities, boxes, numberOfRemainingRoots = args
		# print(args)

		plt.cla() # clear axis
		originalContour.plot(linecolor='k', linestyle='--')
		originalContour.sizePlot()
		for box, numberOfEnclosedRoots in boxes:
			if not hasattr(box, '_color'):
				cmap = plt.get_cmap('jet')
				box._color = cmap(np.random.random())

			plt.text(box.centralPoint.real, box.centralPoint.imag, numberOfEnclosedRoots)
			box.plot(linecolor=box._color)

		plt.scatter(np.real(roots), np.imag(roots))

		rootsLabel = ax.text(0.02, 0.95, 'Zeros remaining: %i'%numberOfRemainingRoots, transform=ax.transAxes)

		fig.canvas.draw()

	if returnAnim:
		return animation.FuncAnimation(fig, update_frame, frames=list(rootFinder), interval=500, repeat_delay=2000)

	elif automaticAnimation:
		ani = animation.FuncAnimation(fig, update_frame, frames=rootFinder, interval=500)

	else:
		def draw_next(event):
			if event.key == ' ':
				update_frame(next(rootFinder))

		fig.canvas.mpl_connect('key_press_event', draw_next)

	plt.show()

	return roots, multiplicities

def showRoots(originalContour, f, df=None, **kwargs):
	"""
	Plots all roots of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.
	"""
	import matplotlib.pyplot as plt
	originalContour.plot(linecolor='k', linestyle='--')
	roots = findRoots(originalContour, f, df, **kwargs)
	plt.scatter(np.real(roots), np.imag(roots), color='k', marker='x')
	plt.show()

<commit_msg>Small box warning<commit_after>from __future__ import division
import numpy as np
import cmath
from scipy import pi, exp, sin
import scipy
from numpy.random import uniform
import itertools
from collections import deque
import warnings

from .IterativeMethods import iterateToRoot
from .CountRoots import prod, RootError

def subdivide(boxDeque, parentBox, parentBox_numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax):
	numberOfRoots = None
	for subBoxes in parentBox.subdivisions():
		try:
			numberOfRoots = [box.count_roots(f, df, integerTol, integrandUpperBound, divMax) for box in np.array(subBoxes)]
			if parentBox_numberOfRoots == sum(numberOfRoots):
				break
		except RootError:
			# If the number of zeros within either of the new contours is not an integer then it is
			# likely that the introduced line which subdivides 'parentBox' lies on a zero.
			# To avoid this we will try to place the subdividing line at a different point along
			# the division axis
			continue

	if numberOfRoots is None or parentBox_numberOfRoots != sum(numberOfRoots):
		# The list of subdivisions has been exhaused and still the number of enclosed zeros does not add up
		raise RuntimeError("""Unable to subdivide box:
			\t%s
			Consider increasing the integrandUpperBound to allow contours closer to roots to be integrated.""" % parentBox)

	boxDeque.extend([(box, numberOfRoots[i]) for i, box in enumerate(subBoxes) if numberOfRoots[i] != 0])

def addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter, multiplicity=None):
	# check that the root we have found is distinct from the ones we already have
	if not roots or np.all(abs(np.array(roots) - root) > newtonStepTol):
		# add the root to the list if it is within the original box
		if originalContour.contains(root):
			roots.append(root)
			if multiplicity is None:
				multiplicities.append(find_multiplicity(root, f, df, rootErrTol, dx=newtonStepTol))
			else:
				multiplicities.append(multiplicity)

		# check to see if there are any other roots implied by the given symmetry
		if guessRootSymmetry is not None:
			for x0 in guessRootSymmetry(root):
				root = iterateToRoot(x0, f, df, newtonStepTol, rootErrTol, newtonMaxIter)
				if root is not None:
					addRoot(root, roots, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

def find_multiplicity(root, f, df=None, rootErrTol=1e-12, dx=1e-8):
	# given a root of the function f find determine the multiplicity
	# of the root numerically
	i, f_root = 0, f(root)
	while abs(f_root) < rootErrTol:
		i += 1
		if df == None:
			f_root = scipy.misc.derivative(f, root, dx=dx, n=i, order=2*i+1)
		else:
			if i == 1:
				f_root = df(root)
			else:
				f_root = scipy.misc.derivative(df, root, dx=dx, n=i-1, order=2*(i-1)+1)
	return i

def findRootsGen(originalContour, f, df=None, guessRoot=[], guessRootSymmetry=None,
	newtonStepTol=1e-8, newtonMaxIter=20, rootErrTol=1e-12,
	absTol=1e-12, relTol=1e-12, divMax=20, integerTol=0.25, integrandUpperBound=1e3,
	M=5):
	"""
	A generator which at each step takes a contour and either finds
	all the zeros of f within it or subdivides it further.  Based
	on the algorithm in [KB]

	Parameters
	----------
	originalContour : subclass of Contour
		The contour C which bounds the region in which all the
		roots of f(z) are sought.
	f : function
		A function of a single complex variable, z, which is
		analytic within C and has no poles or roots on C.
	df : function, optional
		A function of a single complex variable which is the
		derivative of the function f(z). If df is not given
		then it will be approximated with a finite difference
		formula.
	guessRoot : list, optional
		A list of (root, multiplicity) tuples giving the already
		known roots of the function f.
	guessRootSymmetry : function, optional
		A function of a single complex variable, z, which
		returns a list of all points which are expected to be
		roots of f, given that z is a root of f.  It is assumed
		that the symmetric roots have the same multiplicity of
		the original root.
	newtonStepTol : float, optional
		The iterative method used to give a final value for each
		root will exit if the step size, dx, between sucessive
		iterations satisfied abs(dx) < newtonStepTol
	newtonMaxIter : int, optional
		The iterative method used to give a final value for each
		root will exit if the number of iterations exceeds newtonMaxIter.
	rootErrTol : float, optional
		The iterative method used to give a final value for each
		root will exit if, for a point approimating the root, x,
		abs(f(x)) < rootErrTol.
	absTol : float, optional
		The integration along a path will return a result if the
		difference between the last two iterations is less than abstol.
	relTol : float, optional
		The integration along a path will return a result if the
		difference between the last two iterations are less than
		relTol multiplied by the last iteration.
	divMax : int, optional
		The maximum number of divisions before the Romberg integration
		routine of a path exits.
	integerTol : float, optional
		If the result of a contour integration is expected to be
		an integer (ie. when computing the number of zeros within
		a contour) then the result of the last iteration of the
		intergration must be within integerTol of an integer.
		The parameters absTol and relTol are not used in this case.
	integrandUpperBound : float, optional
		The maximum allowed value of abs(df(z)/f(z)) along a contour.
		If abs(df(z)/f(z)) is found to exceed this value then there
		is likely to be a root close to the contour and the
		contour is discarded and a new subdivision created.
		If integrandUpperBound is too large then integrals may take
		a very long time to converge and it is generally be more
		efficient to allow the rootfinding procedure to instead
		choose another contour then spend time evaluting the
		integral along a contour very close to a root.
	M : int, optional
		If the number of roots (including multiplicites) within a
		contour is greater than M then the contour is subdivided
		further.  M must be greater than or equal to the largest
		multiplcity of any root.

	Yields
	------
	list
		Roots of f(z) within the contour C
	list
		Multiplicites of roots
	deque
		The contours which still contain roots
	int
		Remaining number of roots to be found within the contour

	References
	----------
	[KB] Peter Kravanja, Marc Van Barel, "Computing the Zeros of
	Anayltic Functions", Springer (2000)
	"""
	try:
		# total number of zeros, including multiplicities
		totNumberOfRoots = originalContour.count_roots(f, df, integerTol, integrandUpperBound, divMax)
	except RuntimeError:
		raise RuntimeError("""
			Integration along the initial contour has failed.  There is likely a root on or close to the initial contour
			Try either changing the initial contour, if possible, or increasing the integrandUpperBound to allow for
			a longer integration time.""")

	smallBoxWarning = False
	roots = []
	multiplicities = []

	# Add given roots
	# XXX: check these roots and multiplcities
	for root, multiplicity in guessRoot:
		addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

	print('Tot number of Roots', totNumberOfRoots)

	boxes = deque()
	boxes.append((originalContour,totNumberOfRoots))
	while boxes:
		box, numberOfRoots = boxes.pop()

		# if box is smaller than the newtonStepTol then just assume that the root is
		# at the center of the box, print a warning and move on
		if box.area < newtonStepTol:
			smallBoxWarning = True
			if smallBoxWarning is False:
				warnings.warn('The area of the interior of a contour containing %i is smaller than newtonStepTol!  \
					\nThe center of the box has been recored as a root of multiplicity %i but this could not be verified.  \
					\nThe same assumption will be made for future contours this small without an additional warning.'%(numberOfRoots, numberOfRoots))
			root = box.centralPoint
			addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter, multiplicity=numberOfRoots)
			continue

		if numberOfRoots > M:
			subdivide(boxes, box, numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)

		else:
			# approximate the roots in this box
			approxRoots, approxRootMultiplicities = box.approximate_roots(f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)
			for approxRoot, multiplicity in list(zip(approxRoots, approxRootMultiplicities)):
				# if all the roots within the box have been located then just exit
				if numberOfRoots == sum([int(round(multiplicity.real)) for root, multiplicity in zip(roots, multiplicities) if box.contains(root)]):
					break

				if abs(f(approxRoot)) < rootErrTol:
					# the approximate root is good enough
					root = approxRoot
				else:
					# attempt to refine the root
					root = iterateToRoot(approxRoot, f, df, newtonStepTol, rootErrTol, newtonMaxIter)

				if root is not None:
					# if we found a root add it to the list of known roots
					addRoot(root, roots, multiplicities, originalContour, f, df, guessRootSymmetry, newtonStepTol, rootErrTol, newtonMaxIter)

			# if we haven't found all the roots then subdivide further
			if numberOfRoots != sum([int(round(multiplicity.real)) for root, multiplicity in zip(roots, multiplicities) if box.contains(root)]):
				subdivide(boxes, box, numberOfRoots, f, df, absTol, relTol, integerTol, integrandUpperBound, divMax)

		yield roots, multiplicities, boxes, totNumberOfRoots - len(roots)

	if totNumberOfRoots == 0:
		yield [], [], deque(), 0

def findRoots(originalContour, f, df=None, **kwargs):
	"""
	Return a list of all roots of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.
	"""
	rootFinder = findRootsGen(originalContour, f, df, **kwargs)
	for roots, multiplicities, boxes, numberOfRemainingRoots in rootFinder:
		pass
	return roots, multiplicities

def demo_findRoots(originalContour, f, df=None, automaticAnimation=False, returnAnim=False, **kwargs):
	"""
	An interactive demonstration of the processess used to find all the roots
	of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.

	If automaticAnimation is False (default) then press the SPACE key
	to step the animation forward.

	If automaticAnimation is True then the animation will play automatically
	until all the roots have been found.

	If returnAnim is true the animating object returned by matplotlib's animation.FuncAnimation
	will be returned, rather than the animation be shown.
	"""
	import matplotlib.pyplot as plt
	from matplotlib import animation
	rootFinder = findRootsGen(originalContour, f, df, **kwargs)

	originalContour.plot(linecolor='k', linestyle='--')

	fig = plt.gcf()
	ax = plt.gca()

	def update_frame(args):
		roots, multiplicities, boxes, numberOfRemainingRoots = args
		# print(args)

		plt.cla() # clear axis
		originalContour.plot(linecolor='k', linestyle='--')
		originalContour.sizePlot()
		for box, numberOfEnclosedRoots in boxes:
			if not hasattr(box, '_color'):
				cmap = plt.get_cmap('jet')
				box._color = cmap(np.random.random())

			plt.text(box.centralPoint.real, box.centralPoint.imag, numberOfEnclosedRoots)
			box.plot(linecolor=box._color)

		plt.scatter(np.real(roots), np.imag(roots))

		rootsLabel = ax.text(0.02, 0.95, 'Zeros remaining: %i'%numberOfRemainingRoots, transform=ax.transAxes)

		fig.canvas.draw()

	if returnAnim:
		return animation.FuncAnimation(fig, update_frame, frames=list(rootFinder), interval=500, repeat_delay=2000)

	elif automaticAnimation:
		ani = animation.FuncAnimation(fig, update_frame, frames=rootFinder, interval=500)

	else:
		def draw_next(event):
			if event.key == ' ':
				update_frame(next(rootFinder))

		fig.canvas.mpl_connect('key_press_event', draw_next)

	plt.show()

	return roots, multiplicities

def showRoots(originalContour, f, df=None, **kwargs):
	"""
	Plots all roots of a given function f within a given originalContour.
	Shares key word arguments with :func:`cxroots.RootFinder.findRootsGen`.
	"""
	import matplotlib.pyplot as plt
	originalContour.plot(linecolor='k', linestyle='--')
	roots = findRoots(originalContour, f, df, **kwargs)
	plt.scatter(np.real(roots), np.imag(roots), color='k', marker='x')
	plt.show()

<|endoftext|>


<commit_before>"""
Main API for Vega-lite spec generation.

DSL mapping Vega types to IPython traitlets.
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import (
    parse_shorthand, infer_vegalite_type,
    sanitize_dataframe, dataframe_to_json
)
from ._py3k_compat import string_types

import pandas as pd

from . import schema

from .schema import AggregateOp
from .schema import AxisConfig
from .schema import AxisOrient
from .schema import CellConfig
from .schema import Config
from .schema import DataFormat
from .schema import FacetConfig
from .schema import FacetGridConfig
from .schema import FacetScaleConfig
from .schema import FontStyle
from .schema import FontWeight
from .schema import HorizontalAlign
from .schema import LegendConfig
from .schema import MarkConfig
from .schema import NiceTime
from .schema import Scale
from .schema import ScaleConfig
from .schema import SortField
from .schema import SortOrder
from .schema import StackOffset
from .schema import TimeUnit
from .schema import Transform
from .schema import VerticalAlign
from .schema import Data

from .utils import INV_TYPECODE_MAP, TYPE_ABBR

#*************************************************************************
# Channels
#*************************************************************************

class _ChannelMixin(object):

    def _infer_type(self, data):
        if isinstance(data, pd.DataFrame):
            if not self.type and self.field in data:
                self.type = infer_vegalite_type(data[self.field])
        if data is None:
            self.type = ''

    def to_dict(self):
        if not self.field:
            return None
        if not self.type:
            raise ValueError("No vegalite data type defined for {0}".format(self.field))
        return super(_ChannelMixin, self).to_dict()

class PositionChannelDef(_ChannelMixin, schema.PositionChannelDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(PositionChannelDef, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class X(PositionChannelDef):
    channel_name = 'x'

class Y(PositionChannelDef):
    channel_name = 'y'

class Row(PositionChannelDef):
    channel_name = 'row'

class Column(PositionChannelDef):
    channel_name = 'column'

class ChannelDefWithLegend(_ChannelMixin, schema.ChannelDefWithLegend):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(ChannelDefWithLegend, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Color(ChannelDefWithLegend):
    channel_name = 'color'

class Size(ChannelDefWithLegend):
    channel_name = 'size'

class Shape(ChannelDefWithLegend):
    channel_name = 'shape'

class Field(_ChannelMixin, schema.FieldDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(Field, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Text(Field):
    channel_name = 'text'

class Label(Field):
    channel_name = 'label'

class Detail(Field):
    channel_name = 'detail'

class OrderChannel(_ChannelMixin, schema.OrderChannelDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(OrderChannel, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Order(OrderChannel):
    channel_name = 'detail'

class Path(OrderChannel):
    channel_name = 'path'

#*************************************************************************
# Aliases
#*************************************************************************

class Axis(schema.AxisProperties):
    pass

class Bin(schema.BinProperties):
    pass

class Legend(schema.LegendProperties):
    pass

class Formula(schema.VgFormula):

    def __init__(self, field, **kwargs):
        kwargs['field'] = field
        super(Formula, self).__init__(**kwargs)

#*************************************************************************
# Encoding
#*************************************************************************

CHANNEL_CLASSES = {
    'x': X,
    'y': Y,
    'row': Row,
    'column': Column,
    'color': Color,
    'size': Size,
    'shape': Shape,
    'detail': Detail,
    'text': Text,
    'label': Label,
    'path': Path,
    'order': Order

}

CHANNEL_NAMES = list(CHANNEL_CLASSES.keys())

MARK_TYPES = [
    "area",
    "bar",
    "line",
    "point",
    "text",
    "tick",
    "circle",
    "square"
]

class Encoding(schema.Encoding):

    # Position channels
    x = T.Union([T.Instance(X), T.Unicode()],
                default_value=None, allow_none=True)
    y = T.Union([T.Instance(Y), T.Unicode()],
                default_value=None, allow_none=True)
    row = T.Union([T.Instance(Row), T.Unicode()],
                  default_value=None, allow_none=True)
    column = T.Union([T.Instance(Column), T.Unicode()],
                  default_value=None, allow_none=True)

    # Channels with legends
    color = T.Union([T.Instance(Color), T.Unicode()],
                    default_value=None, allow_none=True)
    size = T.Union([T.Instance(Size), T.Unicode()],
                   default_value=None, allow_none=True)
    shape = T.Union([T.Instance(Shape), T.Unicode()],
                    default_value=None, allow_none=True)

    # Field and order channels
    detail = T.Union([T.Instance(Detail), T.List(T.Instance(Detail)), T.Unicode()],
                     default_value=None, allow_none=True)
    text = T.Union([T.Instance(Text), T.Unicode()],
                   default_value=None, allow_none=True)
    label = T.Union([T.Instance(Label), T.Unicode()],
                   default_value=None, allow_none=True)
    path = T.Union([T.Instance(Path), T.List(T.Instance(Path)), T.Unicode()],
                   default_value=None, allow_none=True)
    order = T.Union([T.Instance(Order), T.List(T.Instance(Order)), T.Unicode()],
                   default_value=None, allow_none=True)

    parent = T.Instance(schema.BaseObject, default_value=None, allow_none=True)

    skip = ['parent']

    def _infer_types(self, data):
        for attr in CHANNEL_NAMES:
            val = getattr(self, attr)
            if val is not None:
                val._infer_type(data)

    @T.observe(*CHANNEL_NAMES)
    def _channel_changed(self, change):
        new = change['new']
        name = change['name']
        klass = CHANNEL_CLASSES[name]
        if isinstance(new, string_types):
            setattr(name, klass(new))
        if isinstance(self.parent, Layer):
            channel = getattr(self, name, None)
            if channel is not None:
                meth = getattr(channel, '_infer_type', None)
                if meth is not None:
                    meth(self.parent.data)

#*************************************************************************
# Encoding
#*************************************************************************

class Layer(schema.BaseObject):

    _data = None

    name = T.Unicode()
    description = T.Unicode()
    transform = T.Instance(schema.Transform, default_value=None, allow_none=True)
    mark = T.Enum(MARK_TYPES, default_value='point')
    encoding = T.Instance(Encoding, default_value=None, allow_none=True)
    config = T.Instance(schema.Config, allow_none=True)

    def _encoding_changed(self, name, old, new):
        if isinstance(new, Encoding):
            self.encoding.parent = self
            if isinstance(self.data, pd.DataFrame):
                self.encoding._infer_types(self.data)

    def _set_data(self, new):
        if not (isinstance(new, pd.DataFrame) or isinstance(new, Data) or new is None):
            raise TypeError('Expected DataFrame or altair.Data, got: {0}'.format(new))
        if self.encoding is not None:
            self.encoding._infer_types(new)
        self._data = new

    def _get_data(self):
        return self._data

    data = property(_get_data, _set_data)

    skip = ['data', '_data']

    def __init__(self, *args, **kwargs):
        super(Layer, self).__init__(**kwargs)
        if len(args)==1:
            self.data = args[0]

    def to_dict(self):
        D = super(Layer, self).to_dict()
        if isinstance(self.data, Data):
            D['data'] = self.data.to_dict()
        if isinstance(self.data, pd.DataFrame):
            values = sanitize_dataframe(self.data).to_dict(orient='records')
            D['data'] = Data(values=values).to_dict()
        return D

    def encode(self, *args, **kwargs):
        """Define the encoding for the Layer."""
        for item in args:
            if item.channel_name in kwargs:
                raise ValueError('Mulitple value for {0} provided'.format(item.channel_name))
            kwargs[item.channel_name] = item
        self.encoding = Encoding(**kwargs)
        return self

    def configure(self, **kwargs):
        """Set chart configuration"""
        self.config = schema.Config(**kwargs)
        return self

    def area(self):
        self.mark = 'area'
        return self

    def bar(self):
        self.mark = 'bar'
        return self

    def line(self):
        self.mark = 'line'
        return self

    def point(self):
        self.mark = 'point'
        return self

    def text(self):
        self.mark = 'text'
        return self

    def tick(self):
        self.mark = 'tick'
        return self

    def circle(self):
        self.mark = 'circle'
        return self

    def square(self):
        self.mark = 'square'
        return self

    def _ipython_display_(self):
        from IPython.display import display
        from vega import VegaLite
        display(VegaLite(self.to_dict()))

    def display(self):
        from IPython.display import display
        display(self)

<commit_msg>make API slightly cleaner<commit_after>"""
Main API for Vega-lite spec generation.

DSL mapping Vega types to IPython traitlets.
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import (
    parse_shorthand, infer_vegalite_type,
    sanitize_dataframe, dataframe_to_json
)
from ._py3k_compat import string_types

import pandas as pd

from . import schema

from .schema import AggregateOp
from .schema import AxisConfig
from .schema import AxisOrient
from .schema import CellConfig
from .schema import Config
from .schema import DataFormat
from .schema import FacetConfig
from .schema import FacetGridConfig
from .schema import FacetScaleConfig
from .schema import FontStyle
from .schema import FontWeight
from .schema import HorizontalAlign
from .schema import LegendConfig
from .schema import MarkConfig
from .schema import NiceTime
from .schema import Scale
from .schema import ScaleConfig
from .schema import SortField
from .schema import SortOrder
from .schema import StackOffset
from .schema import TimeUnit
from .schema import Transform
from .schema import VerticalAlign
from .schema import Data

from .utils import INV_TYPECODE_MAP, TYPE_ABBR

#*************************************************************************
# Channels
#*************************************************************************

class _ChannelMixin(object):

    def _infer_type(self, data):
        if isinstance(data, pd.DataFrame):
            if not self.type and self.field in data:
                self.type = infer_vegalite_type(data[self.field])
        if data is None:
            self.type = ''

    def to_dict(self):
        if not self.field:
            return None
        if not self.type:
            raise ValueError("No vegalite data type defined for {0}".format(self.field))
        return super(_ChannelMixin, self).to_dict()

class PositionChannelDef(_ChannelMixin, schema.PositionChannelDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(PositionChannelDef, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class X(PositionChannelDef):
    channel_name = 'x'

class Y(PositionChannelDef):
    channel_name = 'y'

class Row(PositionChannelDef):
    channel_name = 'row'

class Column(PositionChannelDef):
    channel_name = 'column'

class ChannelDefWithLegend(_ChannelMixin, schema.ChannelDefWithLegend):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(ChannelDefWithLegend, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Color(ChannelDefWithLegend):
    channel_name = 'color'

class Size(ChannelDefWithLegend):
    channel_name = 'size'

class Shape(ChannelDefWithLegend):
    channel_name = 'shape'

class Field(_ChannelMixin, schema.FieldDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(Field, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Text(Field):
    channel_name = 'text'

class Label(Field):
    channel_name = 'label'

class Detail(Field):
    channel_name = 'detail'

class OrderChannel(_ChannelMixin, schema.OrderChannelDef):

    skip = ['shorthand']

    shorthand = T.Unicode('')
    type = T.Union([schema.Type(), T.Unicode()],
                   allow_none=True, default_value=None)

    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(OrderChannel, self).__init__(**kwargs)

    @T.observe('shorthand')
    def _shorthand_changed(self, change):
        D = parse_shorthand(change['new'])
        for key, val in D.items():
            setattr(self, key, val)

    @T.observe('type')
    def _type_changed(self, change):
        new = change['new']
        if new in TYPE_ABBR:
            self.type = INV_TYPECODE_MAP[new]

class Order(OrderChannel):
    channel_name = 'order'

class Path(OrderChannel):
    channel_name = 'path'

#*************************************************************************
# Aliases
#*************************************************************************

class Axis(schema.AxisProperties):
    pass

class Bin(schema.BinProperties):
    pass

class Legend(schema.LegendProperties):
    pass

class Formula(schema.VgFormula):

    def __init__(self, field, **kwargs):
        kwargs['field'] = field
        super(Formula, self).__init__(**kwargs)

#*************************************************************************
# Encoding
#*************************************************************************

CHANNEL_CLASSES = {
    'x': X,
    'y': Y,
    'row': Row,
    'column': Column,
    'color': Color,
    'size': Size,
    'shape': Shape,
    'detail': Detail,
    'text': Text,
    'label': Label,
    'path': Path,
    'order': Order

}

CHANNEL_NAMES = list(CHANNEL_CLASSES.keys())

MARK_TYPES = [
    "area",
    "bar",
    "line",
    "point",
    "text",
    "tick",
    "circle",
    "square"
]

class Encoding(schema.Encoding):

    # Position channels
    x = T.Union([T.Instance(X), T.Unicode()],
                default_value=None, allow_none=True)
    y = T.Union([T.Instance(Y), T.Unicode()],
                default_value=None, allow_none=True)
    row = T.Union([T.Instance(Row), T.Unicode()],
                  default_value=None, allow_none=True)
    column = T.Union([T.Instance(Column), T.Unicode()],
                  default_value=None, allow_none=True)

    # Channels with legends
    color = T.Union([T.Instance(Color), T.Unicode()],
                    default_value=None, allow_none=True)
    size = T.Union([T.Instance(Size), T.Unicode()],
                   default_value=None, allow_none=True)
    shape = T.Union([T.Instance(Shape), T.Unicode()],
                    default_value=None, allow_none=True)

    # Field and order channels
    detail = T.Union([T.Instance(Detail), T.List(T.Instance(Detail)), T.Unicode()],
                     default_value=None, allow_none=True)
    text = T.Union([T.Instance(Text), T.Unicode()],
                   default_value=None, allow_none=True)
    label = T.Union([T.Instance(Label), T.Unicode()],
                   default_value=None, allow_none=True)
    path = T.Union([T.Instance(Path), T.List(T.Instance(Path)), T.Unicode()],
                   default_value=None, allow_none=True)
    order = T.Union([T.Instance(Order), T.List(T.Instance(Order)), T.Unicode()],
                   default_value=None, allow_none=True)

    parent = T.Instance(schema.BaseObject, default_value=None, allow_none=True)

    skip = ['parent']

    def _infer_types(self, data):
        for attr in CHANNEL_NAMES:
            val = getattr(self, attr)
            if val is not None:
                val._infer_type(data)

    @T.observe(*CHANNEL_NAMES)
    def _channel_changed(self, change):
        new = change['new']
        name = change['name']
        klass = CHANNEL_CLASSES[name]
        if isinstance(new, string_types):
            setattr(name, klass(new))
        if isinstance(self.parent, Layer):
            channel = getattr(self, name, None)
            if channel is not None:
                meth = getattr(channel, '_infer_type', None)
                if meth is not None:
                    meth(self.parent.data)

#*************************************************************************
# Encoding
#*************************************************************************

class Layer(schema.BaseObject):

    _data = None

    name = T.Unicode()
    description = T.Unicode()
    transform = T.Instance(schema.Transform, default_value=None, allow_none=True)
    mark = T.Enum(MARK_TYPES, default_value='point')
    encoding = T.Instance(Encoding, default_value=None, allow_none=True)
    config = T.Instance(schema.Config, allow_none=True)

    def _encoding_changed(self, name, old, new):
        if isinstance(new, Encoding):
            self.encoding.parent = self
            if isinstance(self.data, pd.DataFrame):
                self.encoding._infer_types(self.data)

    def _set_data(self, new):
        if not (isinstance(new, pd.DataFrame) or isinstance(new, Data) or new is None):
            raise TypeError('Expected DataFrame or altair.Data, got: {0}'.format(new))
        if self.encoding is not None:
            self.encoding._infer_types(new)
        self._data = new

    def _get_data(self):
        return self._data

    data = property(_get_data, _set_data)

    skip = ['data', '_data']

    def __init__(self, *args, **kwargs):
        super(Layer, self).__init__(**kwargs)
        if len(args)==1:
            self.data = args[0]

    def to_dict(self):
        D = super(Layer, self).to_dict()
        if isinstance(self.data, Data):
            D['data'] = self.data.to_dict()
        if isinstance(self.data, pd.DataFrame):
            values = sanitize_dataframe(self.data).to_dict(orient='records')
            D['data'] = Data(values=values).to_dict()
        return D

    def encode(self, *args, **kwargs):
        """Define the encoding for the Layer."""
        for item in args:
            if item.channel_name in kwargs:
                raise ValueError('Mulitple value for {0} provided'.format(item.channel_name))
            kwargs[item.channel_name] = item
        for key, val in kwargs.items():
            if not isinstance(val, CHANNEL_CLASSES[key]):
                kwargs[key] = CHANNEL_CLASSES[key](val)
        self.encoding = Encoding(**kwargs)
        return self

    def configure(self, **kwargs):
        """Set chart configuration"""
        self.config = schema.Config(**kwargs)
        return self

    def area(self):
        self.mark = 'area'
        return self

    def bar(self):
        self.mark = 'bar'
        return self

    def line(self):
        self.mark = 'line'
        return self

    def point(self):
        self.mark = 'point'
        return self

    def text(self):
        self.mark = 'text'
        return self

    def tick(self):
        self.mark = 'tick'
        return self

    def circle(self):
        self.mark = 'circle'
        return self

    def square(self):
        self.mark = 'square'
        return self

    def _ipython_display_(self):
        from IPython.display import display
        from vega import VegaLite
        display(VegaLite(self.to_dict()))

    def display(self):
        from IPython.display import display
        display(self)

<|endoftext|>