<commit_before># Copyright 2018 The Cirq Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import shutil
import sys
from typing import Optional, Iterable, Callable, cast

from dev_tools import shell_tools, git_env_tools
from dev_tools.github_repository import GithubRepository
from dev_tools.prepared_env import PreparedEnv

def get_unhidden_ungenerated_python_files(directory: str) -> Iterable[str]:
    """Iterates through relevant python files within the given directory.

    Args:
        directory: The top-level directory to explore.

    Yields:
        File paths.
    """
    for dirpath, dirnames, filenames in os.walk(directory, topdown=True):
        if os.path.split(dirpath)[-1].startswith('.'):
            dirnames.clear()
            continue

        for filename in filenames:
            if filename.endswith('.py') and not filename.endswith('_pb2.py'):
                yield os.path.join(dirpath, filename)

def create_virtual_env(venv_path: str,
                       requirements_paths: Iterable[str],
                       python_path: str,
                       verbose: bool) -> None:
    """Creates a new virtual environment and then installs dependencies.

    Args:
        venv_path: Where to put the virtual environment's state.
        requirements_paths: Location of requirements files to -r install.
        python_path: The python binary to use.
        verbose: When set, more progress output is produced.
    """
    shell_tools.run_cmd('virtualenv',
                        None if verbose else '--quiet',
                        '-p',
                        python_path,
                        venv_path,
                        out=sys.stderr)
    pip_path = os.path.join(venv_path, 'bin', 'pip')
    for req_path in requirements_paths:
        shell_tools.run_cmd(pip_path,
                            'install',
                            None if verbose else '--quiet',
                            '-r',
                            req_path,
                            out=sys.stderr)

def prepare_temporary_test_environment(
        destination_directory: str,
        repository: GithubRepository,
        pull_request_number: Optional[int],
        verbose: bool,
        env_name: str = '.test_virtualenv',
        python_path: str = '/usr/bin/python3.5',
        commit_ids_known_callback: Callable[[PreparedEnv], None] = None
) -> PreparedEnv:
    """Prepares a temporary test environment at the (existing empty) directory.

    Args:
        destination_directory: The location to put files. The caller is
            responsible for deleting the directory, whether or not this method
             succeeds or fails.
        repository: The github repository to download content from, if a pull
            request number is given.
        pull_request_number: If set, test content is fetched from github.
            Otherwise copies of local files are used.
        verbose: When set, more progress output is produced.
        env_name: The name to use for the virtual environment.
        python_path: Location of the python binary to use within the
            virtual environment.
        commit_ids_known_callback: A function to call when the actual commit id
            being tested is known, before the virtual environment is ready.

    Returns:
        Commit ids corresponding to content to test/compare.
    """
    # Fetch content.
    if pull_request_number is not None:
        env = git_env_tools.fetch_github_pull_request(
            destination_directory=destination_directory,
            repository=repository,
            pull_request_number=pull_request_number,
            verbose=verbose)
    else:
        env = git_env_tools.fetch_local_files(
            destination_directory=destination_directory,
            verbose=verbose)

    if commit_ids_known_callback is not None:
        commit_ids_known_callback(env)

    # Create virtual environment.
    base_path = cast(str, env.destination_directory)
    env_path = os.path.join(base_path, env_name)
    req_path = os.path.join(base_path, 'requirements.txt')
    req_path_2 = os.path.join(base_path,
                              'dev_tools',
                              'conf',
                              'pip-list-dev-tools.txt')
    create_virtual_env(venv_path=env_path,
                       python_path=python_path,
                       requirements_paths=[req_path, req_path_2],
                       verbose=verbose)

    return PreparedEnv(github_repo=env.repository,
                       actual_commit_id=env.actual_commit_id,
                       compare_commit_id=env.compare_commit_id,
                       destination_directory=env.destination_directory,
                       virtual_env_path=env_path)

def derive_temporary_python2_environment(
        destination_directory: str,
        python3_environment: PreparedEnv,
        verbose: bool,
        env_name: str = '.test_virtualenv_py2',
        python_path: str = "/usr/bin/python2.7") -> PreparedEnv:
    """Creates a python 2.7 environment starting from a prepared python 3 one.

    Args:
        destination_directory: Where to put the python 2 environment.
        python3_environment: The prepared environment to start from.
        verbose: When set, more progress output is produced.
        env_name: The name to use for the virtualenv directory.
        python_path: The python binary to use.

    Returns:
        A description of the environment that was prepared.
    """

    shutil.rmtree(destination_directory)
    input_directory = cast(str, python3_environment.destination_directory)
    os.chdir(input_directory)
    conversion_script_path = os.path.join(
        input_directory,
        'dev_tools',
        'python2.7-generate.sh')
    shell_tools.run_cmd('bash',
                        conversion_script_path,
                        destination_directory,
                        input_directory,
                        python3_environment.virtual_env_path,
                        out=sys.stderr)
    os.chdir(destination_directory)

    # Create virtual environment.
    env_path = os.path.join(destination_directory, env_name)
    # (These files are output by dev_tools/python2.7-generate.sh.)
    req_path = os.path.join(destination_directory, 'requirements.txt')
    req_path_2 = os.path.join(destination_directory, 'pip-list-test-tools.txt')
    create_virtual_env(venv_path=env_path,
                       python_path=python_path,
                       requirements_paths=[req_path, req_path_2],
                       verbose=verbose)

    return PreparedEnv(github_repo=python3_environment.repository,
                       actual_commit_id=python3_environment.actual_commit_id,
                       compare_commit_id=python3_environment.compare_commit_id,
                       destination_directory=destination_directory,
                       virtual_env_path=env_path)
<commit_msg>Add contrib requirements to preparing test environment (#1170)<commit_after># Copyright 2018 The Cirq Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import shutil
import sys
from typing import Optional, Iterable, Callable, cast

from dev_tools import shell_tools, git_env_tools
from dev_tools.github_repository import GithubRepository
from dev_tools.prepared_env import PreparedEnv

def get_unhidden_ungenerated_python_files(directory: str) -> Iterable[str]:
    """Iterates through relevant python files within the given directory.

    Args:
        directory: The top-level directory to explore.

    Yields:
        File paths.
    """
    for dirpath, dirnames, filenames in os.walk(directory, topdown=True):
        if os.path.split(dirpath)[-1].startswith('.'):
            dirnames.clear()
            continue

        for filename in filenames:
            if filename.endswith('.py') and not filename.endswith('_pb2.py'):
                yield os.path.join(dirpath, filename)

def create_virtual_env(venv_path: str,
                       requirements_paths: Iterable[str],
                       python_path: str,
                       verbose: bool) -> None:
    """Creates a new virtual environment and then installs dependencies.

    Args:
        venv_path: Where to put the virtual environment's state.
        requirements_paths: Location of requirements files to -r install.
        python_path: The python binary to use.
        verbose: When set, more progress output is produced.
    """
    shell_tools.run_cmd('virtualenv',
                        None if verbose else '--quiet',
                        '-p',
                        python_path,
                        venv_path,
                        out=sys.stderr)
    pip_path = os.path.join(venv_path, 'bin', 'pip')
    for req_path in requirements_paths:
        shell_tools.run_cmd(pip_path,
                            'install',
                            None if verbose else '--quiet',
                            '-r',
                            req_path,
                            out=sys.stderr)

def prepare_temporary_test_environment(
        destination_directory: str,
        repository: GithubRepository,
        pull_request_number: Optional[int],
        verbose: bool,
        env_name: str = '.test_virtualenv',
        python_path: str = '/usr/bin/python3.5',
        commit_ids_known_callback: Callable[[PreparedEnv], None] = None
) -> PreparedEnv:
    """Prepares a temporary test environment at the (existing empty) directory.

    Args:
        destination_directory: The location to put files. The caller is
            responsible for deleting the directory, whether or not this method
             succeeds or fails.
        repository: The github repository to download content from, if a pull
            request number is given.
        pull_request_number: If set, test content is fetched from github.
            Otherwise copies of local files are used.
        verbose: When set, more progress output is produced.
        env_name: The name to use for the virtual environment.
        python_path: Location of the python binary to use within the
            virtual environment.
        commit_ids_known_callback: A function to call when the actual commit id
            being tested is known, before the virtual environment is ready.

    Returns:
        Commit ids corresponding to content to test/compare.
    """
    # Fetch content.
    if pull_request_number is not None:
        env = git_env_tools.fetch_github_pull_request(
            destination_directory=destination_directory,
            repository=repository,
            pull_request_number=pull_request_number,
            verbose=verbose)
    else:
        env = git_env_tools.fetch_local_files(
            destination_directory=destination_directory,
            verbose=verbose)

    if commit_ids_known_callback is not None:
        commit_ids_known_callback(env)

    # Create virtual environment.
    base_path = cast(str, env.destination_directory)
    env_path = os.path.join(base_path, env_name)
    req_path = os.path.join(base_path, 'requirements.txt')
    dev_req_path = os.path.join(base_path,
                                'dev_tools',
                                'conf',
                                'pip-list-dev-tools.txt')
    contrib_req_path = os.path.join(base_path,
                                    'cirq',
                                    'contrib',
                                    'contrib-requirements.txt')
    rev_paths = [req_path, dev_req_path, contrib_req_path]
    create_virtual_env(venv_path=env_path,
                       python_path=python_path,
                       requirements_paths=rev_paths,
                       verbose=verbose)

    return PreparedEnv(github_repo=env.repository,
                       actual_commit_id=env.actual_commit_id,
                       compare_commit_id=env.compare_commit_id,
                       destination_directory=env.destination_directory,
                       virtual_env_path=env_path)

def derive_temporary_python2_environment(
        destination_directory: str,
        python3_environment: PreparedEnv,
        verbose: bool,
        env_name: str = '.test_virtualenv_py2',
        python_path: str = "/usr/bin/python2.7") -> PreparedEnv:
    """Creates a python 2.7 environment starting from a prepared python 3 one.

    Args:
        destination_directory: Where to put the python 2 environment.
        python3_environment: The prepared environment to start from.
        verbose: When set, more progress output is produced.
        env_name: The name to use for the virtualenv directory.
        python_path: The python binary to use.

    Returns:
        A description of the environment that was prepared.
    """

    shutil.rmtree(destination_directory)
    input_directory = cast(str, python3_environment.destination_directory)
    os.chdir(input_directory)
    conversion_script_path = os.path.join(
        input_directory,
        'dev_tools',
        'python2.7-generate.sh')
    shell_tools.run_cmd('bash',
                        conversion_script_path,
                        destination_directory,
                        input_directory,
                        python3_environment.virtual_env_path,
                        out=sys.stderr)
    os.chdir(destination_directory)

    # Create virtual environment.
    env_path = os.path.join(destination_directory, env_name)
    # (These files are output by dev_tools/python2.7-generate.sh.)
    req_path = os.path.join(destination_directory, 'requirements.txt')
    dev_req_path = os.path.join(destination_directory,
                                'pip-list-test-tools.txt')
    contrib_req_path = os.path.join(destination_directory,
                                    'cirq',
                                    'contrib',
                                    'contrib-requirements.txt')
    req_paths = [req_path, dev_req_path, contrib_req_path]
    create_virtual_env(venv_path=env_path,
                       python_path=python_path,
                       requirements_paths=req_paths,
                       verbose=verbose)

    return PreparedEnv(github_repo=python3_environment.repository,
                       actual_commit_id=python3_environment.actual_commit_id,
                       compare_commit_id=python3_environment.compare_commit_id,
                       destination_directory=destination_directory,
                       virtual_env_path=env_path)
<|endoftext|>


<commit_before># -*- Mode: Python; indent-tabs-mode: nil -*-

# verify the entire block chain.

from caesure.script import parse_script, pprint_script, ScriptError, OPCODES
from caesure._script import ScriptError
from caesure.block_db import BlockDB
from caesure.proto import Name
from caesure.ledger import LedgerState
from caesure.ansi import *

import sys
W = sys.stderr.write

def get_names (db, name, h, stop):
    names = []
    while 1:
        if h == stop:
            break
        names.append ((name, h))
        name = db.prev[name]
        h -= 1
    names.reverse()
    return names

fails = open ('fails.txt', 'wb')

# wrap verify to catch all failures, while pretending they didn't happen.
from caesure.bitcoin import TX
original_verify = TX.verify

def verify_wrapper (self, index, lock_script, block_timestamp):
    try:
        original_verify (self, index, lock_script, block_timestamp)
    except:
        W ('%064x %d failed\n' % (self.name, index))
        fails.write (
            '%s %s %d %s %d\n' % (
                hex(self.name),
                self.raw.encode ('hex'),
                index,
                lock_script.encode ('hex'),
                block_timestamp
            )
        )

TX.verify = verify_wrapper

def main (db, G):
    h, name = db.get_highest_uncontested_block()
    names = get_names (db, name, h, -1)
    WR ('scanning %d blocks...\n' % (len(names),))
    lx = LedgerState (load=False)
    lx.do_yields = False
    height = 0
    for name, height in names:
        #W ('name=%r, height=%d\n' % (repr(name), height))
        b = db[name]
        #W ('b=%r\n' % (b,))
        lx.feed_block (b, height, verify=True)
        if height % 1000 == 0:
            WB ('[%d]' % (height,))

if __name__ == '__main__':
    import argparse
    class GlobalState:
        pass
    G = GlobalState()
    p = argparse.ArgumentParser()
    p.add_argument ('-b', '--base', help='data directory', default='/usr/local/caesure', metavar='PATH')
    p.add_argument ('--stop', type=int, help='stopping block height', default=10000)
    G.args = p.parse_args()
    db = G.block_db = BlockDB (read_only=True)
    #import coro.backdoor
    #coro.spawn (coro.backdoor.serve, unix_path='/tmp/verify.bd')
    #coro.spawn (main, db, G)
    #coro.event_loop()
    main (db, G)

<commit_msg>verify_wrapper: format output for test_txns.py + commented tx.name.<commit_after># -*- Mode: Python; indent-tabs-mode: nil -*-

# verify the entire block chain.

from caesure.script import parse_script, pprint_script, ScriptError, OPCODES
from caesure._script import ScriptError
from caesure.block_db import BlockDB
from caesure.proto import Name
from caesure.ledger import LedgerState
from caesure.ansi import *

import sys
W = sys.stderr.write

def get_names (db, name, h, stop):
    names = []
    while 1:
        if h == stop:
            break
        names.append ((name, h))
        name = db.prev[name]
        h -= 1
    names.reverse()
    return names

fails = open ('fails.txt', 'wb')

# wrap verify to catch all failures, while pretending they didn't happen.
from caesure.bitcoin import TX
original_verify = TX.verify

def verify_wrapper (self, index, lock_script, block_timestamp):
    try:
        original_verify (self, index, lock_script, block_timestamp)
    except:
        W ('%064x %d failed\n' % (self.name, index))
        fails.write (
            '#%s\n%s %s %d %d\n' % (
                hex(self.name),
                lock_script.encode ('hex'),
                self.raw.encode ('hex'),
                index,
                block_timestamp
            )
        )

TX.verify = verify_wrapper

def main (db, G):
    h, name = db.get_highest_uncontested_block()
    names = get_names (db, name, h, -1)
    WR ('scanning %d blocks...\n' % (len(names),))
    lx = LedgerState (load=False)
    lx.do_yields = False
    height = 0
    for name, height in names:
        #W ('name=%r, height=%d\n' % (repr(name), height))
        b = db[name]
        #W ('b=%r\n' % (b,))
        lx.feed_block (b, height, verify=True)
        if height % 1000 == 0:
            WB ('[%d]' % (height,))

if __name__ == '__main__':
    import argparse
    class GlobalState:
        pass
    G = GlobalState()
    p = argparse.ArgumentParser()
    p.add_argument ('-b', '--base', help='data directory', default='/usr/local/caesure', metavar='PATH')
    p.add_argument ('--stop', type=int, help='stopping block height', default=10000)
    G.args = p.parse_args()
    db = G.block_db = BlockDB (read_only=True)
    #import coro.backdoor
    #coro.spawn (coro.backdoor.serve, unix_path='/tmp/verify.bd')
    #coro.spawn (main, db, G)
    #coro.event_loop()
    main (db, G)

<|endoftext|>


<commit_before>import pytest
from storage.kv_store_leveldb import KeyValueStorageLeveldb
from storage.kv_store_rocksdb import KeyValueStorageRocksdb
from storage.kv_store import KeyValueStorage

i = 0

@pytest.yield_fixture(scope="function", params=['rocksdb', 'leveldb'])
def kv(request, tempdir) -> KeyValueStorage:
    global i

    if request.param == 'leveldb':
        kv = KeyValueStorageLeveldb(tempdir, 'kv{}'.format(i))
    else:
        kv = KeyValueStorageRocksdb(tempdir, 'kv{}'.format(i))

    assert kv.read_only == False

    kv.put('k1', 'v1')
    kv.put('k2', 'v2')
    kv.put('k3', 'v3')
    kv.close()

    if request.param == 'leveldb':
        kv = KeyValueStorageLeveldb(tempdir, 'kv{}'.format(i), read_only=True)
    else:
        kv = KeyValueStorageRocksdb(tempdir, 'kv{}'.format(i), read_only=True)

    i += 1
    yield kv
    kv.close()

def test_read_only_get(kv):
    assert kv.read_only == True

    v = kv.get('k1')
    assert b'v1' == v

    v = kv.get('k2')
    assert b'v2' == v

    v = kv.get('k3')
    assert b'v3' == v

def test_read_only_put(kv):
    assert kv.read_only == True

    with pytest.raises(Exception, match="Not supported operation in read only mode."):
        kv.put('k4', 'v4')

def test_read_only_remove(kv):
    assert kv.read_only == True

    with pytest.raises(Exception, match="Not supported operation in read only mode."):
        kv.remove('k1')
<commit_msg>Fix flake8.<commit_after>import pytest
from storage.kv_store_leveldb import KeyValueStorageLeveldb
from storage.kv_store_rocksdb import KeyValueStorageRocksdb
from storage.kv_store import KeyValueStorage

i = 0

@pytest.yield_fixture(scope="function", params=['rocksdb', 'leveldb'])
def kv(request, tempdir) -> KeyValueStorage:
    global i

    if request.param == 'leveldb':
        kv = KeyValueStorageLeveldb(tempdir, 'kv{}'.format(i))
    else:
        kv = KeyValueStorageRocksdb(tempdir, 'kv{}'.format(i))

    assert kv.read_only is False

    kv.put('k1', 'v1')
    kv.put('k2', 'v2')
    kv.put('k3', 'v3')
    kv.close()

    if request.param == 'leveldb':
        kv = KeyValueStorageLeveldb(tempdir, 'kv{}'.format(i), read_only=True)
    else:
        kv = KeyValueStorageRocksdb(tempdir, 'kv{}'.format(i), read_only=True)

    i += 1
    yield kv
    kv.close()

def test_read_only_get(kv):
    assert kv.read_only is True

    v = kv.get('k1')
    assert b'v1' == v

    v = kv.get('k2')
    assert b'v2' == v

    v = kv.get('k3')
    assert b'v3' == v

def test_read_only_put(kv):
    assert kv.read_only is True

    with pytest.raises(Exception, match="Not supported operation in read only mode."):
        kv.put('k4', 'v4')

def test_read_only_remove(kv):
    assert kv.read_only is True

    with pytest.raises(Exception, match="Not supported operation in read only mode."):
        kv.remove('k1')
<|endoftext|>


<commit_before># SMOP -- Simple Matlab/Octave to Python compiler
# Copyright 2011-2016 Victor Leikehman

from __future__ import print_function

import py_compile
import tempfile
import fnmatch
import tarfile
import sys
import os
import traceback
from os.path import basename, splitext

import lexer
import options
import parse
import resolve
import backend
from version import __version__

def print_header(fp):
    if options.no_header:
        return
    print("# Generated with SMOP ", __version__, file=fp)
    print("from libsmop import *", file=fp)
    print("#", options.filename, file=fp)

def main():
    if "M" in options.debug:
        import pdb
        pdb.set_trace()
    if not options.filelist:
        options.parser.print_help()
        return
    if options.output == "-":
        fp = sys.stdout
    elif options.output:
        fp = open(options.output, "w")
    else:
        fp = None
    if fp:
        print_header(fp)

    nerrors = 0
    for i, options.filename in enumerate(options.filelist):
        try:
            if options.verbose:
                print(i, options.filename)
            if not options.filename.endswith(".m"):
                print("\tIgnored: '%s' (unexpected file type)" %
                      options.filename)
                continue
            if basename(options.filename) in options.xfiles:
                if options.verbose:
                    print("\tExcluded: '%s'" % options.filename)
                continue
            buf = open(options.filename).read()
            buf = buf.replace("\r\n", "\n")
            # FIXME buf = buf.decode("ascii", errors="ignore")
            stmt_list = parse.parse(buf if buf[-1] == '\n' else buf + '\n')

            if not stmt_list:
                continue
            if not options.no_resolve:
                G = resolve.resolve(stmt_list)
            if not options.no_backend:
                s = backend.backend(stmt_list)
            if not options.output:
                f = splitext(basename(options.filename))[0] + ".py"
                with open(f, "w") as fp:
                    print_header(fp)
                    fp.write(s)
            else:
                fp.write(s)
        except KeyboardInterrupt:
            break
        except:
            nerrors += 1
            traceback.print_exc(file=sys.stdout)
            if options.strict:
                break
        finally:
            pass
    if nerrors:
        print("Errors:", nerrors)

if __name__ == "__main__":
    main()
<commit_msg>I think missing version.py is fixed now<commit_after># SMOP -- Simple Matlab/Octave to Python compiler
# Copyright 2011-2016 Victor Leikehman

from __future__ import print_function

import py_compile
import tempfile
import fnmatch
import tarfile
import sys
import os
import traceback
from os.path import basename, splitext

from . import options
from . import parse
from . import resolve
from . import backend
from version import __version__

def print_header(fp):
    if options.no_header:
        return
    print("# Generated with SMOP ", __version__, file=fp)
    print("from libsmop import *", file=fp)
    print("#", options.filename, file=fp)

def main():
    if "M" in options.debug:
        import pdb
        pdb.set_trace()
    if not options.filelist:
        options.parser.print_help()
        return
    if options.output == "-":
        fp = sys.stdout
    elif options.output:
        fp = open(options.output, "w")
    else:
        fp = None
    if fp:
        print_header(fp)

    nerrors = 0
    for i, options.filename in enumerate(options.filelist):
        try:
            if options.verbose:
                print(i, options.filename)
            if not options.filename.endswith(".m"):
                print("\tIgnored: '%s' (unexpected file type)" %
                      options.filename)
                continue
            if basename(options.filename) in options.xfiles:
                if options.verbose:
                    print("\tExcluded: '%s'" % options.filename)
                continue
            buf = open(options.filename).read()
            buf = buf.replace("\r\n", "\n")
            # FIXME buf = buf.decode("ascii", errors="ignore")
            stmt_list = parse.parse(buf if buf[-1] == '\n' else buf + '\n')

            if not stmt_list:
                continue
            if not options.no_resolve:
                G = resolve.resolve(stmt_list)
            if not options.no_backend:
                s = backend.backend(stmt_list)
            if not options.output:
                f = splitext(basename(options.filename))[0] + ".py"
                with open(f, "w") as fp:
                    print_header(fp)
                    fp.write(s)
            else:
                fp.write(s)
        except KeyboardInterrupt:
            break
        except:
            nerrors += 1
            traceback.print_exc(file=sys.stdout)
            if options.strict:
                break
        finally:
            pass
    if nerrors:
        print("Errors:", nerrors)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>from collections import OrderedDict
from cached_property import cached_property

import numpy as np
from sympy import Dummy
from sympy.core.decorators import call_highest_priority
from sympy.core.sympify import converter as sympify_converter

from devito.finite_differences import Differentiable
from devito.types.basic import AbstractTensor
from devito.types.dense import Function, TimeFunction
from devito.types.utils import NODE

__all__ = ['TensorFunction', 'TensorTimeFunction', 'VectorFunction', 'VectorTimeFunction']

class TensorFunction(AbstractTensor):
    """
    Tensor valued Function represented as a Matrix.
    Each component is a Function or TimeFunction.

    A TensorFunction and the classes that inherit from it takes the same parameters as
    a DiscreteFunction and additionally:

    Parameters
    ----------
    name : str
        Name of the symbol.
    grid : Grid, optional
        Carries shape, dimensions, and dtype of the TensorFunction. When grid is not
        provided, shape and dimensions must be given. For MPI execution, a
        Grid is compulsory.
    space_order : int or 3-tuple of ints, optional
        Discretisation order for space derivatives. Defaults to 1. ``space_order`` also
        impacts the number of points available around a generic point of interest.  By
        default, ``space_order`` points are available on both sides of a generic point of
        interest, including those nearby the grid boundary. Sometimes, fewer points
        suffice; in other scenarios, more points are necessary. In such cases, instead of
        an integer, one can pass a 3-tuple ``(o, lp, rp)`` indicating the discretization
        order (``o``) as well as the number of points on the left (``lp``) and right
        (``rp``) sides of a generic point of interest.
    shape : tuple of ints, optional
        Shape of the domain region in grid points. Only necessary if ``grid`` isn't given.
    dimensions : tuple of Dimension, optional
        Dimensions associated with the object. Only necessary if ``grid`` isn't given.
    dtype : data-type, optional
        Any object that can be interpreted as a numpy data type. Defaults
        to ``np.float32``.
    staggered : Dimension or tuple of Dimension or Stagger, optional
        Define how the TensorFunction is staggered.
    initializer : callable or any object exposing the buffer interface, optional
        Data initializer. If a callable is provided, data is allocated lazily.
    allocator : MemoryAllocator, optional
        Controller for memory allocation. To be used, for example, when one wants
        to take advantage of the memory hierarchy in a NUMA architecture. Refer to
        `default_allocator.__doc__` for more information.
    padding : int or tuple of ints, optional
        .. deprecated:: shouldn't be used; padding is now automatically inserted.
        Allocate extra grid points to maximize data access alignment. When a tuple
        of ints, one int per Dimension should be provided.
    symmetric : bool, optional
        Whether the tensor is symmetric or not. Defaults to True.
    diagonal : Bool, optional
        Whether the tensor is diagonal or not. Defaults to False.
    staggered: tuple of Dimension, optional
        Staggering of each component, needs to have the size of the tensor. Defaults
        to the Dimensions.
    """

    _is_TimeDependent = False

    _sub_type = Function

    _class_priority = 10
    _op_priority = Differentiable._op_priority + 1.

    def __init_finalize__(self, *args, **kwargs):
        grid = kwargs.get('grid')
        dimensions = kwargs.get('dimensions')
        inds, _ = Function.__indices_setup__(grid=grid,
                                             dimensions=dimensions)
        self._space_dimensions = inds

    @classmethod
    def __subfunc_setup__(cls, *args, **kwargs):
        """
        Creates the components of the TensorFunction
        either from input or from input Dimensions.
        """
        comps = kwargs.get("components")
        if comps is not None:
            return comps

        grid = kwargs.get("grid")
        if grid is None:
            dims = kwargs.get('dimensions')
            if dims is None:
                raise TypeError("Need either `grid` or `dimensions`")
        else:
            dims = grid.dimensions
        stagg = kwargs.get("staggered", None)
        name = kwargs.get("name")
        symm = kwargs.get('symmetric', True)
        diag = kwargs.get('diagonal', False)

        funcs = []
        # Fill tensor, only upper diagonal if symmetric
        for i, d in enumerate(dims):
            funcs2 = [0 for _ in range(len(dims))]
            start = i if (symm or diag) else 0
            stop = i + 1 if diag else len(dims)
            for j in range(start, stop):
                kwargs["name"] = "%s_%s%s" % (name, d.name, dims[j].name)
                kwargs["staggered"] = (stagg[i][j] if stagg is not None
                                       else (NODE if i == j else (d, dims[j])))
                funcs2[j] = cls._sub_type(**kwargs)
            funcs.append(funcs2)

        # Symmetrize and fill diagonal if symmetric
        if symm:
            funcs = np.array(funcs) + np.triu(np.array(funcs), k=1).T
            funcs = funcs.tolist()
        return funcs

    @call_highest_priority('__mul__')
    def __rmul__(self, other):
        """
        Prevents Functions to be interpreted as matrices in 2D becaue of `.shape`
        TODO: Remove after sympy 1.7
        """
        if getattr(other, 'is_DiscreteFunction', False):
            tmp_func = Dummy(other.name)
            simplemul = super(TensorFunction, self).__rmul__(tmp_func)
            return simplemul.subs(tmp_func, other)
        # honest sympy matrices defer to their class's routine
        if getattr(other, 'is_Matrix', False):
            return self._eval_matrix_rmul(other)
        return super(TensorFunction, self).__rmul__(other)

    @call_highest_priority('__rmul__')
    def __mul__(self, other):
        """
        Prevents Functions to be interpreted as matrices in 2D becaue of `.shape`
        TODO: Remove after sympy 1.7
        """
        if getattr(other, 'is_DiscreteFunction', False):
            tmp_func = Dummy(other.name)
            simplemul = super(TensorFunction, self).__mul__(tmp_func)
            return simplemul.subs(tmp_func, other)
        return super(TensorFunction, self).__mul__(other)

    def __getattr__(self, name):
        """
        Try calling a dynamically created FD shortcut.

        Notes
        -----
        This method acts as a fallback for __getattribute__
        """
        try:
            return self.applyfunc(lambda x: x if x == 0 else getattr(x, name))
        except:
            raise AttributeError("%r object has no attribute %r" % (self.__class__, name))

    def _eval_at(self, func):
        """
        Evaluate tensor at func location
        """
        def entries(i, j, func):
            return getattr(self[i, j], '_eval_at', lambda x: self[i, j])(func[i, j])
        entry = lambda i, j: entries(i, j, func)
        return self._new(self.rows, self.cols, entry)

    @classmethod
    def __indices_setup__(cls, **kwargs):
        return Function.__indices_setup__(grid=kwargs.get('grid'),
                                          dimensions=kwargs.get('dimensions'))

    @property
    def _symbolic_functions(self):
        return frozenset().union(*[a._symbolic_functions for a in self.values()])

    @property
    def is_TimeDependent(self):
        return self._is_TimeDependent

    @property
    def space_dimensions(self):
        """Spatial dimensions."""
        return self._space_dimensions

    @cached_property
    def space_order(self):
        """The space order for all components."""
        return ({a.space_order for a in self} - {None}).pop()

    @property
    def is_diagonal(self):
        """Whether the tensor is diagonal."""
        return np.all([self[i, j] == 0 for j in range(self.cols)
                       for i in range(self.rows) if i != j])

    @property
    def evaluate(self):
        return self.applyfunc(lambda x: getattr(x, 'evaluate', x))

    def values(self):
        if self.is_diagonal:
            return [self[i, i] for i in range(self.shape[0])]
        elif self.is_symmetric:
            val = super(TensorFunction, self).values()
            return list(OrderedDict.fromkeys(val))
        else:
            return super(TensorFunction, self).values()

    def div(self, shift=None):
        """
        Divergence of the TensorFunction (is a VectorFunction).
        """
        comps = []
        func = vec_func(self, self)
        for j, d in enumerate(self.space_dimensions):
            comps.append(sum([getattr(self[j, i], 'd%s' % d.name)
                              for i, d in enumerate(self.space_dimensions)]))
        return func._new(comps)

    @property
    def laplace(self):
        """
        Laplacian of the TensorFunction.
        """
        comps = []
        func = vec_func(self, self)
        for j, d in enumerate(self.space_dimensions):
            comps.append(sum([getattr(self[j, i], 'd%s2' % d.name)
                              for i, d in enumerate(self.space_dimensions)]))
        return func._new(comps)

    def grad(self, shift=None):
        raise AttributeError("Gradient of a second order tensor not supported")

    def new_from_mat(self, mat):
        func = tens_func(self, self)
        return func._new(self.rows, self.cols, mat)

    def classof_prod(self, other, mat):
        try:
            is_mat = len(mat[0]) > 1
        except TypeError:
            is_mat = False
        is_time = (getattr(self, '_is_TimeDependent', False) or
                   getattr(other, '_is_TimeDependent', False))
        return mat_time_dict[(is_time, is_mat)]

class VectorFunction(TensorFunction):
    """
    Vector valued space varying Function as a rank 1 tensor of Function.
    """

    is_VectorValued = True
    is_TensorValued = False

    _sub_type = Function

    _is_TimeDependent = False

    @property
    def is_transposed(self):
        return self.shape[0] == 1

    @classmethod
    def __subfunc_setup__(cls, *args, **kwargs):
        """
        Creates the components of the VectorFunction
        either from input or from input dimensions.
        """
        comps = kwargs.get("components")
        if comps is not None:
            return comps
        funcs = []
        grid = kwargs.get("grid")
        if grid is None:
            dims = kwargs.get('dimensions')
            if dims is None:
                raise TypeError("Need either `grid` or `dimensions`")
        else:
            dims = grid.dimensions
        stagg = kwargs.get("staggered", None)
        name = kwargs.get("name")
        for i, d in enumerate(dims):
            kwargs["name"] = "%s_%s" % (name, d.name)
            kwargs["staggered"] = stagg[i] if stagg is not None else d
            funcs.append(cls._sub_type(**kwargs))

        return funcs

    # Custom repr and str
    def __str__(self):
        st = ''.join([' %-2s,' % c for c in self])[1:-1]
        return "Vector(%s)" % st

    __repr__ = __str__

    def div(self, shift=None):
        """
        Divergence of the VectorFunction, creates the divergence Function.
        """
        return sum([getattr(self[i], 'd%s' % d.name)
                    for i, d in enumerate(self.space_dimensions)])

    @property
    def laplace(self):
        """
        Laplacian of the VectorFunction, creates the Laplacian VectorFunction.
        """
        func = vec_func(self, self)
        comps = [sum([getattr(s, 'd%s2' % d.name) for d in self.space_dimensions])
                 for s in self]
        return func._new(comps)

    @property
    def curl(self):
        """
        Gradient of the (3D) VectorFunction, creates the curl VectorFunction.
        """

        if len(self.space_dimensions) != 3:
            raise AttributeError("Curl only supported for 3D VectorFunction")
        # The curl of a VectorFunction is a VectorFunction
        derivs = ['d%s' % d.name for d in self.space_dimensions]
        comp1 = getattr(self[2], derivs[1]) - getattr(self[1], derivs[2])
        comp2 = getattr(self[0], derivs[2]) - getattr(self[2], derivs[0])
        comp3 = getattr(self[1], derivs[0]) - getattr(self[0], derivs[1])

        vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
        return vec_func._new(3, 1, [comp1, comp2, comp3])

    def grad(self, shift=None):
        """
        Gradient of the VectorFunction, creates the gradient TensorFunction.
        """
        func = tens_func(self, self)
        comps = [[getattr(f, 'd%s' % d.name) for d in self.space_dimensions]
                 for f in self]
        return func._new(comps)

    def outer(self, other):
        comps = [[self[i] * other[j] for i in range(self.cols)] for j in range(self.cols)]
        func = tens_func(self, other)
        return func._new(comps)

    def new_from_mat(self, mat):
        func = vec_func(self, self)
        return func._new(self.rows, 1, mat)

class TensorTimeFunction(TensorFunction):
    """
    Time varying TensorFunction.
    """

    is_TensorValued = True

    _sub_type = TimeFunction

    _is_TimeDependent = True

    @cached_property
    def time_order(self):
        """The time order for all components."""
        return ({a.time_order for a in self} - {None}).pop()

class VectorTimeFunction(VectorFunction, TensorTimeFunction):
    """
    Time varying VectorFunction.
    """

    is_VectorValued = True
    is_TensorValued = False

    _sub_type = TimeFunction

    _is_TimeDependent = True
    _time_position = 0

mat_time_dict = {(True, True): TensorTimeFunction, (True, False): VectorTimeFunction,
                 (False, True): TensorFunction, (False, False): VectorFunction}

def vec_func(func1, func2):
    f1 = getattr(func1, 'is_TimeDependent', False)
    f2 = getattr(func2, 'is_TimeDependent', False)
    return VectorTimeFunction if f1 or f2 else VectorFunction

def tens_func(func1, func2):
    f1 = getattr(func1, 'is_TimeDependent', False)
    f2 = getattr(func2, 'is_TimeDependent', False)
    return TensorTimeFunction if f1 or f2 else TensorFunction

def sympify_tensor(arg):
    return arg

sympify_converter[TensorFunction] = sympify_tensor
<commit_msg>FD: add missing shift option<commit_after>from collections import OrderedDict
from cached_property import cached_property

import numpy as np
from sympy import Dummy
from sympy.core.decorators import call_highest_priority
from sympy.core.sympify import converter as sympify_converter

from devito.finite_differences import Differentiable
from devito.types.basic import AbstractTensor
from devito.types.dense import Function, TimeFunction
from devito.types.utils import NODE

__all__ = ['TensorFunction', 'TensorTimeFunction', 'VectorFunction', 'VectorTimeFunction']

class TensorFunction(AbstractTensor):
    """
    Tensor valued Function represented as a Matrix.
    Each component is a Function or TimeFunction.

    A TensorFunction and the classes that inherit from it takes the same parameters as
    a DiscreteFunction and additionally:

    Parameters
    ----------
    name : str
        Name of the symbol.
    grid : Grid, optional
        Carries shape, dimensions, and dtype of the TensorFunction. When grid is not
        provided, shape and dimensions must be given. For MPI execution, a
        Grid is compulsory.
    space_order : int or 3-tuple of ints, optional
        Discretisation order for space derivatives. Defaults to 1. ``space_order`` also
        impacts the number of points available around a generic point of interest.  By
        default, ``space_order`` points are available on both sides of a generic point of
        interest, including those nearby the grid boundary. Sometimes, fewer points
        suffice; in other scenarios, more points are necessary. In such cases, instead of
        an integer, one can pass a 3-tuple ``(o, lp, rp)`` indicating the discretization
        order (``o``) as well as the number of points on the left (``lp``) and right
        (``rp``) sides of a generic point of interest.
    shape : tuple of ints, optional
        Shape of the domain region in grid points. Only necessary if ``grid`` isn't given.
    dimensions : tuple of Dimension, optional
        Dimensions associated with the object. Only necessary if ``grid`` isn't given.
    dtype : data-type, optional
        Any object that can be interpreted as a numpy data type. Defaults
        to ``np.float32``.
    staggered : Dimension or tuple of Dimension or Stagger, optional
        Define how the TensorFunction is staggered.
    initializer : callable or any object exposing the buffer interface, optional
        Data initializer. If a callable is provided, data is allocated lazily.
    allocator : MemoryAllocator, optional
        Controller for memory allocation. To be used, for example, when one wants
        to take advantage of the memory hierarchy in a NUMA architecture. Refer to
        `default_allocator.__doc__` for more information.
    padding : int or tuple of ints, optional
        .. deprecated:: shouldn't be used; padding is now automatically inserted.
        Allocate extra grid points to maximize data access alignment. When a tuple
        of ints, one int per Dimension should be provided.
    symmetric : bool, optional
        Whether the tensor is symmetric or not. Defaults to True.
    diagonal : Bool, optional
        Whether the tensor is diagonal or not. Defaults to False.
    staggered: tuple of Dimension, optional
        Staggering of each component, needs to have the size of the tensor. Defaults
        to the Dimensions.
    """

    _is_TimeDependent = False

    _sub_type = Function

    _class_priority = 10
    _op_priority = Differentiable._op_priority + 1.

    def __init_finalize__(self, *args, **kwargs):
        grid = kwargs.get('grid')
        dimensions = kwargs.get('dimensions')
        inds, _ = Function.__indices_setup__(grid=grid,
                                             dimensions=dimensions)
        self._space_dimensions = inds

    @classmethod
    def __subfunc_setup__(cls, *args, **kwargs):
        """
        Creates the components of the TensorFunction
        either from input or from input Dimensions.
        """
        comps = kwargs.get("components")
        if comps is not None:
            return comps

        grid = kwargs.get("grid")
        if grid is None:
            dims = kwargs.get('dimensions')
            if dims is None:
                raise TypeError("Need either `grid` or `dimensions`")
        else:
            dims = grid.dimensions
        stagg = kwargs.get("staggered", None)
        name = kwargs.get("name")
        symm = kwargs.get('symmetric', True)
        diag = kwargs.get('diagonal', False)

        funcs = []
        # Fill tensor, only upper diagonal if symmetric
        for i, d in enumerate(dims):
            funcs2 = [0 for _ in range(len(dims))]
            start = i if (symm or diag) else 0
            stop = i + 1 if diag else len(dims)
            for j in range(start, stop):
                kwargs["name"] = "%s_%s%s" % (name, d.name, dims[j].name)
                kwargs["staggered"] = (stagg[i][j] if stagg is not None
                                       else (NODE if i == j else (d, dims[j])))
                funcs2[j] = cls._sub_type(**kwargs)
            funcs.append(funcs2)

        # Symmetrize and fill diagonal if symmetric
        if symm:
            funcs = np.array(funcs) + np.triu(np.array(funcs), k=1).T
            funcs = funcs.tolist()
        return funcs

    @call_highest_priority('__mul__')
    def __rmul__(self, other):
        """
        Prevents Functions to be interpreted as matrices in 2D becaue of `.shape`
        TODO: Remove after sympy 1.7
        """
        if getattr(other, 'is_DiscreteFunction', False):
            tmp_func = Dummy(other.name)
            simplemul = super(TensorFunction, self).__rmul__(tmp_func)
            return simplemul.subs(tmp_func, other)
        # honest sympy matrices defer to their class's routine
        if getattr(other, 'is_Matrix', False):
            return self._eval_matrix_rmul(other)
        return super(TensorFunction, self).__rmul__(other)

    @call_highest_priority('__rmul__')
    def __mul__(self, other):
        """
        Prevents Functions to be interpreted as matrices in 2D becaue of `.shape`
        TODO: Remove after sympy 1.7
        """
        if getattr(other, 'is_DiscreteFunction', False):
            tmp_func = Dummy(other.name)
            simplemul = super(TensorFunction, self).__mul__(tmp_func)
            return simplemul.subs(tmp_func, other)
        return super(TensorFunction, self).__mul__(other)

    def __getattr__(self, name):
        """
        Try calling a dynamically created FD shortcut.

        Notes
        -----
        This method acts as a fallback for __getattribute__
        """
        try:
            return self.applyfunc(lambda x: x if x == 0 else getattr(x, name))
        except:
            raise AttributeError("%r object has no attribute %r" % (self.__class__, name))

    def _eval_at(self, func):
        """
        Evaluate tensor at func location
        """
        def entries(i, j, func):
            return getattr(self[i, j], '_eval_at', lambda x: self[i, j])(func[i, j])
        entry = lambda i, j: entries(i, j, func)
        return self._new(self.rows, self.cols, entry)

    @classmethod
    def __indices_setup__(cls, **kwargs):
        return Function.__indices_setup__(grid=kwargs.get('grid'),
                                          dimensions=kwargs.get('dimensions'))

    @property
    def _symbolic_functions(self):
        return frozenset().union(*[a._symbolic_functions for a in self.values()])

    @property
    def is_TimeDependent(self):
        return self._is_TimeDependent

    @property
    def space_dimensions(self):
        """Spatial dimensions."""
        return self._space_dimensions

    @cached_property
    def space_order(self):
        """The space order for all components."""
        return ({a.space_order for a in self} - {None}).pop()

    @property
    def is_diagonal(self):
        """Whether the tensor is diagonal."""
        return np.all([self[i, j] == 0 for j in range(self.cols)
                       for i in range(self.rows) if i != j])

    @property
    def evaluate(self):
        return self.applyfunc(lambda x: getattr(x, 'evaluate', x))

    def values(self):
        if self.is_diagonal:
            return [self[i, i] for i in range(self.shape[0])]
        elif self.is_symmetric:
            val = super(TensorFunction, self).values()
            return list(OrderedDict.fromkeys(val))
        else:
            return super(TensorFunction, self).values()

    def div(self, shift=None):
        """
        Divergence of the TensorFunction (is a VectorFunction).
        """
        comps = []
        func = vec_func(self, self)
        for j, d in enumerate(self.space_dimensions):
            comps.append(sum([getattr(self[j, i], 'd%s' % d.name)
                              for i, d in enumerate(self.space_dimensions)]))
        return func._new(comps)

    @property
    def laplace(self):
        """
        Laplacian of the TensorFunction.
        """
        comps = []
        func = vec_func(self, self)
        for j, d in enumerate(self.space_dimensions):
            comps.append(sum([getattr(self[j, i], 'd%s2' % d.name)
                              for i, d in enumerate(self.space_dimensions)]))
        return func._new(comps)

    def grad(self, shift=None):
        raise AttributeError("Gradient of a second order tensor not supported")

    def new_from_mat(self, mat):
        func = tens_func(self, self)
        return func._new(self.rows, self.cols, mat)

    def classof_prod(self, other, mat):
        try:
            is_mat = len(mat[0]) > 1
        except TypeError:
            is_mat = False
        is_time = (getattr(self, '_is_TimeDependent', False) or
                   getattr(other, '_is_TimeDependent', False))
        return mat_time_dict[(is_time, is_mat)]

class VectorFunction(TensorFunction):
    """
    Vector valued space varying Function as a rank 1 tensor of Function.
    """

    is_VectorValued = True
    is_TensorValued = False

    _sub_type = Function

    _is_TimeDependent = False

    @property
    def is_transposed(self):
        return self.shape[0] == 1

    @classmethod
    def __subfunc_setup__(cls, *args, **kwargs):
        """
        Creates the components of the VectorFunction
        either from input or from input dimensions.
        """
        comps = kwargs.get("components")
        if comps is not None:
            return comps
        funcs = []
        grid = kwargs.get("grid")
        if grid is None:
            dims = kwargs.get('dimensions')
            if dims is None:
                raise TypeError("Need either `grid` or `dimensions`")
        else:
            dims = grid.dimensions
        stagg = kwargs.get("staggered", None)
        name = kwargs.get("name")
        for i, d in enumerate(dims):
            kwargs["name"] = "%s_%s" % (name, d.name)
            kwargs["staggered"] = stagg[i] if stagg is not None else d
            funcs.append(cls._sub_type(**kwargs))

        return funcs

    # Custom repr and str
    def __str__(self):
        st = ''.join([' %-2s,' % c for c in self])[1:-1]
        return "Vector(%s)" % st

    __repr__ = __str__

    def div(self, shift=None):
        """
        Divergence of the VectorFunction, creates the divergence Function.
        """
        return sum([getattr(self[i], 'd%s' % d.name)(x0=None if shift is None else
                                                     d + shift * d.spacing)
                    for i, d in enumerate(self.space_dimensions)])

    @property
    def laplace(self):
        """
        Laplacian of the VectorFunction, creates the Laplacian VectorFunction.
        """
        func = vec_func(self, self)
        comps = [sum([getattr(s, 'd%s2' % d.name) for d in self.space_dimensions])
                 for s in self]
        return func._new(comps)

    @property
    def curl(self):
        """
        Gradient of the (3D) VectorFunction, creates the curl VectorFunction.
        """

        if len(self.space_dimensions) != 3:
            raise AttributeError("Curl only supported for 3D VectorFunction")
        # The curl of a VectorFunction is a VectorFunction
        derivs = ['d%s' % d.name for d in self.space_dimensions]
        comp1 = getattr(self[2], derivs[1]) - getattr(self[1], derivs[2])
        comp2 = getattr(self[0], derivs[2]) - getattr(self[2], derivs[0])
        comp3 = getattr(self[1], derivs[0]) - getattr(self[0], derivs[1])

        vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
        return vec_func._new(3, 1, [comp1, comp2, comp3])

    def grad(self, shift=None):
        """
        Gradient of the VectorFunction, creates the gradient TensorFunction.
        """
        func = tens_func(self, self)
        comps = [[getattr(f, 'd%s' % d.name) for d in self.space_dimensions]
                 for f in self]
        return func._new(comps)

    def outer(self, other):
        comps = [[self[i] * other[j] for i in range(self.cols)] for j in range(self.cols)]
        func = tens_func(self, other)
        return func._new(comps)

    def new_from_mat(self, mat):
        func = vec_func(self, self)
        return func._new(self.rows, 1, mat)

class TensorTimeFunction(TensorFunction):
    """
    Time varying TensorFunction.
    """

    is_TensorValued = True

    _sub_type = TimeFunction

    _is_TimeDependent = True

    @cached_property
    def time_order(self):
        """The time order for all components."""
        return ({a.time_order for a in self} - {None}).pop()

class VectorTimeFunction(VectorFunction, TensorTimeFunction):
    """
    Time varying VectorFunction.
    """

    is_VectorValued = True
    is_TensorValued = False

    _sub_type = TimeFunction

    _is_TimeDependent = True
    _time_position = 0

mat_time_dict = {(True, True): TensorTimeFunction, (True, False): VectorTimeFunction,
                 (False, True): TensorFunction, (False, False): VectorFunction}

def vec_func(func1, func2):
    f1 = getattr(func1, 'is_TimeDependent', False)
    f2 = getattr(func2, 'is_TimeDependent', False)
    return VectorTimeFunction if f1 or f2 else VectorFunction

def tens_func(func1, func2):
    f1 = getattr(func1, 'is_TimeDependent', False)
    f2 = getattr(func2, 'is_TimeDependent', False)
    return TensorTimeFunction if f1 or f2 else TensorFunction

def sympify_tensor(arg):
    return arg

sympify_converter[TensorFunction] = sympify_tensor
<|endoftext|>


<commit_before>from itertools import chain
from devops.helpers import retry

class EnvironmentException(object):
    pass

class ManagedObject(object):
    def __init__(self):
        super(ManagedObject, self).__init__()
        self._driver = None
        self.name = None

    @property
    def driver(self):
        if self._driver is None:
            raise EnvironmentException, "Object '%s' wasn't built yet" % self.name
        return self._driver

    @driver.setter
    def driver(self, driver):
        self._driver = driver

    @driver.deleter
    def driver(self):
        del self._driver

class Environment(ManagedObject):
    def __init__(self, name):
        super(Environment, self).__init__()

        self.name = name
        self.networks = []
        self.nodes = []
        self.built = False

    @property
    def node(self):
        name2node = {}
        for node in self.nodes:
            name2node[node.name] = node
        return name2node

    @property
    def network(self):
        name2network = {}
        for network in self.networks:
            name2network[network.name] = network
        return name2network

class Network(ManagedObject):
    def __init__(self, name, dhcp_server=False, pxe=False,
                 reserve_static=True, forward='nat'):
        super(Network, self).__init__()

        self.name = name
        self.dhcp_server = dhcp_server
        self.reserve_static=reserve_static
        self.interfaces = []
        self.pxe = pxe
        self.forward = forward
        self.environment = None

    def start(self):
        self.driver.start_network(self)

    def stop(self):
        self.driver.stop_network(self)

class Node(ManagedObject):
    def __init__(self, name, cpu=1, memory=512, arch='x86_64', vnc=False):
        super(Node, self).__init__()

        self.name = name

        self.cpu = cpu
        self.memory = memory
        self.arch = arch
        self.vnc = vnc
        self.interfaces = []
        self.bridged_interfaces = []
        self.disks = []
        self.boot = []
        self.cdrom = None
        self.environment = None

    def start(self):
        self.driver.start_node(self)

    def stop(self):
        self.driver.stop_node(self)

    def reset(self):
        self.driver.reset_node(self)

    def reboot(self):
        self.driver.reboot_node(self)

    def shutdown(self):
        self.driver.shutdown_node(self)

    def suspend(self):
        self.driver.suspend_node(self)

    def resume(self):
        self.driver.resume_node(self)

    @property
    def snapshots(self):
        return self.driver.get_node_snapshots(self)

    def has_snapshot(self, snapshot_name):
        return snapshot_name in self.snapshots

    def _delete_and_create(self, name):
        #noinspection PyBroadException
        try:
            self.delete_snapshot(name)
        except:
            pass
        return self.driver.create_snapshot(self, name=name)

    def save_snapshot(self, name=None, force=False):
        "Create node state snapshot. Returns snapshot name"
        if name and force and self.has_snapshot(name):
            return retry(10, self._delete_and_create, name=name)
        return self.driver.create_snapshot(self, name=name)

    def restore_snapshot(self, snapshot_name=None):
        "Revert node state to given snapshot. If no snapshot name given, revert to current snapshot."
        self.driver.revert_snapshot(self, snapshot_name)

    def delete_snapshot(self, snapshot_name=None):
        "Delete snapshot. If no snapshot name given, delete current snapshot."
        self.driver.delete_snapshot(self, snapshot_name)

    def send_keys(self, keys):
        self.driver.send_keys_to_node(self, keys)

    @property
    def ip_addresses(self):
        addresses = []
        for interface in self.interfaces:
            addresses += interface.ip_addresses
        return addresses

    @property
    def ip_address(self):
        x = self.ip_addresses
        if len(x) == 0:
            return None
        return x[0]

    @property
    def ip_address_by_network(self):
        name2ip_addresses = {}
        for interface in self.interfaces:
            if len(interface.ip_addresses) == 0:
                name2ip_addresses[interface.network.name] = None
            else:
                name2ip_addresses[interface.network.name] = interface.ip_addresses[0]
        return name2ip_addresses

    @ManagedObject.driver.setter
    def driver(self, driver):
        ManagedObject.driver.fset(self, driver)
        for interface in self.interfaces:
            interface.driver = driver

class Cdrom(object):
    def __init__(self, isopath=None, bus='ide'):
        self.isopath = isopath
        self.bus = bus

class Disk(object):
    def __init__(self, size=None, path=None, format='qcow2', bus='virtio',
                 base_image=None):
        self.size = size
        self.format = format
        self.bus = bus
        self.path = path
        # Either copy or use "backing file" feature
        self.base_image = base_image

class BridgedInterface(ManagedObject):
    def __init__(self, bridge, type='virtio'):
        super(BridgedInterface, self).__init__()
        self.bridge = bridge
        self.type = type

class Interface(ManagedObject):
    def __init__(self, network, ip_addresses=None, type='virtio'):
        super(Interface, self).__init__()
        if not ip_addresses: ip_addresses = []
        self.node = None
        self.network = network
        if not isinstance(ip_addresses, (list, tuple)):
            ip_addresses = (ip_addresses,)
        self._ip_addresses = ip_addresses
        self.mac_address = None
        self.type = type
        self.target_dev = None

    @property
    def ip_addresses(self):
        return self._ip_addresses

    @ip_addresses.setter
    def ip_addresses(self, value):
        if not isinstance(value, (list, tuple)):
            value = (value,)
        self._ip_addresses = value
<commit_msg>Add metadata field to node<commit_after>from itertools import chain
from devops.helpers import retry

class EnvironmentException(object):
    pass

class ManagedObject(object):
    def __init__(self):
        super(ManagedObject, self).__init__()
        self._driver = None
        self.name = None

    @property
    def driver(self):
        if self._driver is None:
            raise EnvironmentException, "Object '%s' wasn't built yet" % self.name
        return self._driver

    @driver.setter
    def driver(self, driver):
        self._driver = driver

    @driver.deleter
    def driver(self):
        del self._driver

class Environment(ManagedObject):
    def __init__(self, name):
        super(Environment, self).__init__()

        self.name = name
        self.networks = []
        self.nodes = []
        self.built = False

    @property
    def node(self):
        name2node = {}
        for node in self.nodes:
            name2node[node.name] = node
        return name2node

    @property
    def network(self):
        name2network = {}
        for network in self.networks:
            name2network[network.name] = network
        return name2network

class Network(ManagedObject):
    def __init__(self, name, dhcp_server=False, pxe=False,
                 reserve_static=True, forward='nat'):
        super(Network, self).__init__()

        self.name = name
        self.dhcp_server = dhcp_server
        self.reserve_static=reserve_static
        self.interfaces = []
        self.pxe = pxe
        self.forward = forward
        self.environment = None

    def start(self):
        self.driver.start_network(self)

    def stop(self):
        self.driver.stop_network(self)

class Node(ManagedObject):
    def __init__(self, name, cpu=1, memory=512, arch='x86_64', vnc=False,
                 metadata=None):
        super(Node, self).__init__()

        self.name = name
        if not metadata: self.metadata = {}
        self.cpu = cpu
        self.memory = memory
        self.arch = arch
        self.vnc = vnc
        self.interfaces = []
        self.bridged_interfaces = []
        self.disks = []
        self.boot = []
        self.cdrom = None
        self.environment = None

    def start(self):
        self.driver.start_node(self)

    def stop(self):
        self.driver.stop_node(self)

    def reset(self):
        self.driver.reset_node(self)

    def reboot(self):
        self.driver.reboot_node(self)

    def shutdown(self):
        self.driver.shutdown_node(self)

    def suspend(self):
        self.driver.suspend_node(self)

    def resume(self):
        self.driver.resume_node(self)

    @property
    def snapshots(self):
        return self.driver.get_node_snapshots(self)

    def has_snapshot(self, snapshot_name):
        return snapshot_name in self.snapshots

    def _delete_and_create(self, name):
        #noinspection PyBroadException
        try:
            self.delete_snapshot(name)
        except:
            pass
        return self.driver.create_snapshot(self, name=name)

    def save_snapshot(self, name=None, force=False):
        "Create node state snapshot. Returns snapshot name"
        if name and force and self.has_snapshot(name):
            return retry(10, self._delete_and_create, name=name)
        return self.driver.create_snapshot(self, name=name)

    def restore_snapshot(self, snapshot_name=None):
        "Revert node state to given snapshot. If no snapshot name given, revert to current snapshot."
        self.driver.revert_snapshot(self, snapshot_name)

    def delete_snapshot(self, snapshot_name=None):
        "Delete snapshot. If no snapshot name given, delete current snapshot."
        self.driver.delete_snapshot(self, snapshot_name)

    def send_keys(self, keys):
        self.driver.send_keys_to_node(self, keys)

    @property
    def ip_addresses(self):
        addresses = []
        for interface in self.interfaces:
            addresses += interface.ip_addresses
        return addresses

    @property
    def ip_address(self):
        x = self.ip_addresses
        if len(x) == 0:
            return None
        return x[0]

    @property
    def ip_address_by_network(self):
        name2ip_addresses = {}
        for interface in self.interfaces:
            if len(interface.ip_addresses) == 0:
                name2ip_addresses[interface.network.name] = None
            else:
                name2ip_addresses[interface.network.name] = interface.ip_addresses[0]
        return name2ip_addresses

    @ManagedObject.driver.setter
    def driver(self, driver):
        ManagedObject.driver.fset(self, driver)
        for interface in self.interfaces:
            interface.driver = driver

class Cdrom(object):
    def __init__(self, isopath=None, bus='ide'):
        self.isopath = isopath
        self.bus = bus

class Disk(object):
    def __init__(self, size=None, path=None, format='qcow2', bus='virtio',
                 base_image=None):
        self.size = size
        self.format = format
        self.bus = bus
        self.path = path
        # Either copy or use "backing file" feature
        self.base_image = base_image

class BridgedInterface(ManagedObject):
    def __init__(self, bridge, type='virtio'):
        super(BridgedInterface, self).__init__()
        self.bridge = bridge
        self.type = type

class Interface(ManagedObject):
    def __init__(self, network, ip_addresses=None, type='virtio'):
        super(Interface, self).__init__()
        if not ip_addresses: ip_addresses = []
        self.node = None
        self.network = network
        if not isinstance(ip_addresses, (list, tuple)):
            ip_addresses = (ip_addresses,)
        self._ip_addresses = ip_addresses
        self.mac_address = None
        self.type = type
        self.target_dev = None

    @property
    def ip_addresses(self):
        return self._ip_addresses

    @ip_addresses.setter
    def ip_addresses(self, value):
        if not isinstance(value, (list, tuple)):
            value = (value,)
        self._ip_addresses = value
<|endoftext|>


<commit_before>from django.test import TestCase
from django.db import IntegrityError
from django.core.exceptions import ValidationError
from django.db.models import Model
from django.forms import ValidationError, Form
from .models import *
from .forms import AddressField as FormAddressField

class TestForm(Form):
    address = FormAddressField()

class CountryTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.nz = Country.objects.create(name='New Zealand', code='NZ')
        self.be = Country.objects.create(name='Belgium', code='BE')

    def test_ordering(self):
        qs = Country.objects.all()
        self.assertEqual(qs.count(), 3)
        self.assertEqual(qs[0].code, 'AU')
        self.assertEqual(qs[1].code, 'BE')
        self.assertEqual(qs[2].code, 'NZ')

    def test_unique_name(self):
        self.assertRaises(IntegrityError, Country.objects.create, name='Australia', code='**')

    def test_unicode(self):
        self.assertEqual(unicode(self.au), u'Australia')

class StateTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.qld = State.objects.create(name='Queensland', country=self.au)
        self.empty = State.objects.create(country=self.au)
        self.uk = Country.objects.create(name='United Kingdom', code='UK')
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

    def test_required_country(self):
        self.assertRaises(IntegrityError, State.objects.create)

    def test_ordering(self):
        qs = State.objects.all()
        self.assertEqual(qs.count(), 5)
        self.assertEqual(qs[0].name, '')
        self.assertEqual(qs[1].name, 'Queensland')
        self.assertEqual(qs[2].name, 'Tasmania')
        self.assertEqual(qs[3].name, 'Victoria')
        self.assertEqual(qs[4].name, 'Victoria')

    def test_unique_name_country(self):
        State.objects.create(name='Tasmania', country=self.uk)
        self.assertRaises(IntegrityError, State.objects.create, name='Tasmania', country=self.au)

    def test_unicode(self):
        self.assertEqual(unicode(self.vic), u'Victoria, Australia')
        self.assertEqual(unicode(self.empty), u'Australia')

class LocalityTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.uk = Country.objects.create(name='United Kingdom', code='UK')

        self.au_vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.au_tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.au_qld = State.objects.create(name='Queensland', country=self.au)
        self.au_empty = State.objects.create(country=self.au)
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

        self.au_vic_nco = Locality.objects.create(name='Northcote', postal_code='3070', state=self.au_vic)
        self.au_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.au_vic)
        self.au_vic_ftz = Locality.objects.create(name='Fitzroy', state=self.au_vic)
        self.au_vic_empty = Locality.objects.create(state=self.au_vic)
        self.uk_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.uk_vic)

    def test_required_state(self):
        self.assertRaises(IntegrityError, Locality.objects.create)

    def test_ordering(self):
        qs = Locality.objects.all()
        self.assertEqual(qs.count(), 5)
        self.assertEqual(qs[0].name, '')
        self.assertEqual(qs[1].name, 'Fitzroy')
        self.assertEqual(qs[2].name, 'Melbourne')
        self.assertEqual(qs[3].name, 'Northcote')
        self.assertEqual(qs[4].name, 'Melbourne')

    def test_unique_name_state(self):
        Locality.objects.create(name='Melbourne', state=self.au_qld)
        self.assertRaises(IntegrityError, Locality.objects.create, name='Melbourne', state=self.au_vic)

    def test_unicode(self):
        self.assertEqual(unicode(self.au_vic_mel), u'Melbourne, Victoria 3000, Australia')
        self.assertEqual(unicode(self.au_vic_ftz), u'Fitzroy, Victoria, Australia')
        self.assertEqual(unicode(self.au_vic_empty), u'Victoria, Australia')

class AddressTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.uk = Country.objects.create(name='United Kingdom', code='UK')

        self.au_vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.au_tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.au_qld = State.objects.create(name='Queensland', country=self.au)
        self.au_empty = State.objects.create(country=self.au)
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

        self.au_vic_nco = Locality.objects.create(name='Northcote', postal_code='3070', state=self.au_vic)
        self.au_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.au_vic)
        self.au_vic_ftz = Locality.objects.create(name='Fitzroy', state=self.au_vic)
        self.au_vic_empty = Locality.objects.create(state=self.au_vic)
        self.uk_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.uk_vic)

        self.ad1 = Address.objects.create(street_number='1', route='Some Street', locality=self.au_vic_mel,
                                          raw='1 Some Street, Victoria, Melbourne')
        self.ad2 = Address.objects.create(street_number='10', route='Other Street', locality=self.au_vic_mel,
                                          raw='10 Other Street, Victoria, Melbourne')
        self.ad3 = Address.objects.create(street_number='1', route='Some Street', locality=self.au_vic_nco,
                                          raw='1 Some Street, Northcote, Victoria')
        self.ad_empty = Address.objects.create(locality=self.au_vic_nco, raw='Northcote, Victoria')

    def test_required_raw(self):
        obj = Address.objects.create()
        self.assertRaises(ValidationError, obj.clean)

    def test_ordering(self):
        qs = Address.objects.all()
        self.assertEqual(qs.count(), 4)
        self.assertEqual(qs[0].route, 'Other Street')
        self.assertEqual(qs[1].route, 'Some Street')
        self.assertEqual(qs[2].route, '')
        self.assertEqual(qs[3].route, 'Some Street')

    def test_unique_street_address_locality(self):
        Address.objects.create(street_number='10', route='Other Street', locality=self.au_vic_nco)
        self.assertRaises(
            IntegrityError, Address.objects.create,
            street_number='10', route='Other Street', locality=self.au_vic_mel
        )

    def test_unicode(self):
        self.assertEqual(unicode(self.ad1), u'1 Some Street, Melbourne, Victoria 3000, Australia')
        self.assertEqual(unicode(self.ad_empty), u'Northcote, Victoria 3070, Australia')

class AddressFieldTestCase(TestCase):

    class TestModel(Model):
        address = AddressField()

    def setUp(self):
        self.ad1_dict = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'locality': 'Northcote',
            'postal_code': '3070',
            'state': 'Victoria',
            'state_code': 'VIC',
            'country': 'Australia',
            'country_code': 'AU'
        }
        self.test = self.TestModel()

    def test_assignment_from_dict(self):
        self.test.address = self.ad1_dict
        self.assertEqual(self.test.address.raw, self.ad1_dict['raw'])
        self.assertEqual(self.test.address.street_number, self.ad1_dict['street_number'])
        self.assertEqual(self.test.address.route, self.ad1_dict['route'])
        self.assertEqual(self.test.address.locality.name, self.ad1_dict['locality'])
        self.assertEqual(self.test.address.locality.postal_code, self.ad1_dict['postal_code'])
        self.assertEqual(self.test.address.locality.state.name, self.ad1_dict['state'])
        self.assertEqual(self.test.address.locality.state.code, self.ad1_dict['state_code'])
        self.assertEqual(self.test.address.locality.state.country.name, self.ad1_dict['country'])
        self.assertEqual(self.test.address.locality.state.country.code, self.ad1_dict['country_code'])

    def test_assignment_from_dict_no_country(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'locality': 'Northcote',
            'state': 'Victoria',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_no_state(self):
        ad = {
            'raw': 'Somewhere',
            'locality': 'Northcote',
            'country': 'Australia',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_no_locality(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'state': 'Victoria',
            'country': 'Australia',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_only_address(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, ad['street_number'])
        self.assertEqual(self.test.address.route, ad['route'])
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_string(self):
        self.test.address = self.ad1_dict['raw']
        self.assertEqual(self.test.address.raw, self.ad1_dict['raw'])

    def test_save(self):
        self.test.address = self.ad1_dict
        self.test.save()
        test = self.TestModel.objects.all()[0]
        self.assertEqual(test.address.raw, self.ad1_dict['raw'])
        self.assertEqual(test.address.street_number, self.ad1_dict['street_number'])
        self.assertEqual(test.address.route, self.ad1_dict['route'])
        self.assertEqual(test.address.locality.name, self.ad1_dict['locality'])
        self.assertEqual(test.address.locality.postal_code, self.ad1_dict['postal_code'])
        self.assertEqual(test.address.locality.state.name, self.ad1_dict['state'])
        self.assertEqual(test.address.locality.state.code, self.ad1_dict['state_code'])
        self.assertEqual(test.address.locality.state.country.name, self.ad1_dict['country'])
        self.assertEqual(test.address.locality.state.country.code, self.ad1_dict['country_code'])

class FormAddressFieldTestCase(TestCase):

    def setUp(self):
        self.form = TestForm()
        self.field = self.form.base_fields['address']

    def test_to_python_none(self):
        self.assertEqual(self.field.to_python(None), None)

    def test_to_python_empty(self):
        self.assertEqual(self.field.to_python(''), None)

    def test_to_python_invalid_lat_lng(self):
        self.assertRaises(ValidationError, self.field.to_python, {'latitude': 'x'})
        self.assertRaises(ValidationError, self.field.to_python, {'longitude': 'x'})

    def test_to_python(self):
        res = self.field.to_python({'raw': 'Someplace'})
        self.assertEqual(res.raw, 'Someplace')
<commit_msg>Disabling unique unit-test.<commit_after>from django.test import TestCase
from django.db import IntegrityError
from django.core.exceptions import ValidationError
from django.db.models import Model
from django.forms import ValidationError, Form
from .models import *
from .forms import AddressField as FormAddressField

class TestForm(Form):
    address = FormAddressField()

class CountryTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.nz = Country.objects.create(name='New Zealand', code='NZ')
        self.be = Country.objects.create(name='Belgium', code='BE')

    def test_ordering(self):
        qs = Country.objects.all()
        self.assertEqual(qs.count(), 3)
        self.assertEqual(qs[0].code, 'AU')
        self.assertEqual(qs[1].code, 'BE')
        self.assertEqual(qs[2].code, 'NZ')

    def test_unique_name(self):
        self.assertRaises(IntegrityError, Country.objects.create, name='Australia', code='**')

    def test_unicode(self):
        self.assertEqual(unicode(self.au), u'Australia')

class StateTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.qld = State.objects.create(name='Queensland', country=self.au)
        self.empty = State.objects.create(country=self.au)
        self.uk = Country.objects.create(name='United Kingdom', code='UK')
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

    def test_required_country(self):
        self.assertRaises(IntegrityError, State.objects.create)

    def test_ordering(self):
        qs = State.objects.all()
        self.assertEqual(qs.count(), 5)
        self.assertEqual(qs[0].name, '')
        self.assertEqual(qs[1].name, 'Queensland')
        self.assertEqual(qs[2].name, 'Tasmania')
        self.assertEqual(qs[3].name, 'Victoria')
        self.assertEqual(qs[4].name, 'Victoria')

    def test_unique_name_country(self):
        State.objects.create(name='Tasmania', country=self.uk)
        self.assertRaises(IntegrityError, State.objects.create, name='Tasmania', country=self.au)

    def test_unicode(self):
        self.assertEqual(unicode(self.vic), u'Victoria, Australia')
        self.assertEqual(unicode(self.empty), u'Australia')

class LocalityTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.uk = Country.objects.create(name='United Kingdom', code='UK')

        self.au_vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.au_tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.au_qld = State.objects.create(name='Queensland', country=self.au)
        self.au_empty = State.objects.create(country=self.au)
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

        self.au_vic_nco = Locality.objects.create(name='Northcote', postal_code='3070', state=self.au_vic)
        self.au_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.au_vic)
        self.au_vic_ftz = Locality.objects.create(name='Fitzroy', state=self.au_vic)
        self.au_vic_empty = Locality.objects.create(state=self.au_vic)
        self.uk_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.uk_vic)

    def test_required_state(self):
        self.assertRaises(IntegrityError, Locality.objects.create)

    def test_ordering(self):
        qs = Locality.objects.all()
        self.assertEqual(qs.count(), 5)
        self.assertEqual(qs[0].name, '')
        self.assertEqual(qs[1].name, 'Fitzroy')
        self.assertEqual(qs[2].name, 'Melbourne')
        self.assertEqual(qs[3].name, 'Northcote')
        self.assertEqual(qs[4].name, 'Melbourne')

    def test_unique_name_state(self):
        Locality.objects.create(name='Melbourne', state=self.au_qld)
        self.assertRaises(IntegrityError, Locality.objects.create, name='Melbourne', state=self.au_vic)

    def test_unicode(self):
        self.assertEqual(unicode(self.au_vic_mel), u'Melbourne, Victoria 3000, Australia')
        self.assertEqual(unicode(self.au_vic_ftz), u'Fitzroy, Victoria, Australia')
        self.assertEqual(unicode(self.au_vic_empty), u'Victoria, Australia')

class AddressTestCase(TestCase):

    def setUp(self):
        self.au = Country.objects.create(name='Australia', code='AU')
        self.uk = Country.objects.create(name='United Kingdom', code='UK')

        self.au_vic = State.objects.create(name='Victoria', code='VIC', country=self.au)
        self.au_tas = State.objects.create(name='Tasmania', code='TAS', country=self.au)
        self.au_qld = State.objects.create(name='Queensland', country=self.au)
        self.au_empty = State.objects.create(country=self.au)
        self.uk_vic = State.objects.create(name='Victoria', code='VIC', country=self.uk)

        self.au_vic_nco = Locality.objects.create(name='Northcote', postal_code='3070', state=self.au_vic)
        self.au_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.au_vic)
        self.au_vic_ftz = Locality.objects.create(name='Fitzroy', state=self.au_vic)
        self.au_vic_empty = Locality.objects.create(state=self.au_vic)
        self.uk_vic_mel = Locality.objects.create(name='Melbourne', postal_code='3000', state=self.uk_vic)

        self.ad1 = Address.objects.create(street_number='1', route='Some Street', locality=self.au_vic_mel,
                                          raw='1 Some Street, Victoria, Melbourne')
        self.ad2 = Address.objects.create(street_number='10', route='Other Street', locality=self.au_vic_mel,
                                          raw='10 Other Street, Victoria, Melbourne')
        self.ad3 = Address.objects.create(street_number='1', route='Some Street', locality=self.au_vic_nco,
                                          raw='1 Some Street, Northcote, Victoria')
        self.ad_empty = Address.objects.create(locality=self.au_vic_nco, raw='Northcote, Victoria')

    def test_required_raw(self):
        obj = Address.objects.create()
        self.assertRaises(ValidationError, obj.clean)

    def test_ordering(self):
        qs = Address.objects.all()
        self.assertEqual(qs.count(), 4)
        self.assertEqual(qs[0].route, 'Other Street')
        self.assertEqual(qs[1].route, 'Some Street')
        self.assertEqual(qs[2].route, '')
        self.assertEqual(qs[3].route, 'Some Street')

    # def test_unique_street_address_locality(self):
    #     Address.objects.create(street_number='10', route='Other Street', locality=self.au_vic_nco)
    #     self.assertRaises(
    #         IntegrityError, Address.objects.create,
    #         street_number='10', route='Other Street', locality=self.au_vic_mel
    #     )

    def test_unicode(self):
        self.assertEqual(unicode(self.ad1), u'1 Some Street, Melbourne, Victoria 3000, Australia')
        self.assertEqual(unicode(self.ad_empty), u'Northcote, Victoria 3070, Australia')

class AddressFieldTestCase(TestCase):

    class TestModel(Model):
        address = AddressField()

    def setUp(self):
        self.ad1_dict = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'locality': 'Northcote',
            'postal_code': '3070',
            'state': 'Victoria',
            'state_code': 'VIC',
            'country': 'Australia',
            'country_code': 'AU'
        }
        self.test = self.TestModel()

    def test_assignment_from_dict(self):
        self.test.address = self.ad1_dict
        self.assertEqual(self.test.address.raw, self.ad1_dict['raw'])
        self.assertEqual(self.test.address.street_number, self.ad1_dict['street_number'])
        self.assertEqual(self.test.address.route, self.ad1_dict['route'])
        self.assertEqual(self.test.address.locality.name, self.ad1_dict['locality'])
        self.assertEqual(self.test.address.locality.postal_code, self.ad1_dict['postal_code'])
        self.assertEqual(self.test.address.locality.state.name, self.ad1_dict['state'])
        self.assertEqual(self.test.address.locality.state.code, self.ad1_dict['state_code'])
        self.assertEqual(self.test.address.locality.state.country.name, self.ad1_dict['country'])
        self.assertEqual(self.test.address.locality.state.country.code, self.ad1_dict['country_code'])

    def test_assignment_from_dict_no_country(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'locality': 'Northcote',
            'state': 'Victoria',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_no_state(self):
        ad = {
            'raw': 'Somewhere',
            'locality': 'Northcote',
            'country': 'Australia',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_no_locality(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
            'state': 'Victoria',
            'country': 'Australia',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, '')
        self.assertEqual(self.test.address.route, '')
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_dict_only_address(self):
        ad = {
            'raw': '1 Somewhere Street, Northcote, Victoria 3070, VIC, AU',
            'street_number': '1',
            'route': 'Somewhere Street',
        }
        self.test.address = ad
        self.assertEqual(self.test.address.raw, ad['raw'])
        self.assertEqual(self.test.address.street_number, ad['street_number'])
        self.assertEqual(self.test.address.route, ad['route'])
        self.assertEqual(self.test.address.locality, None)

    def test_assignment_from_string(self):
        self.test.address = self.ad1_dict['raw']
        self.assertEqual(self.test.address.raw, self.ad1_dict['raw'])

    def test_save(self):
        self.test.address = self.ad1_dict
        self.test.save()
        test = self.TestModel.objects.all()[0]
        self.assertEqual(test.address.raw, self.ad1_dict['raw'])
        self.assertEqual(test.address.street_number, self.ad1_dict['street_number'])
        self.assertEqual(test.address.route, self.ad1_dict['route'])
        self.assertEqual(test.address.locality.name, self.ad1_dict['locality'])
        self.assertEqual(test.address.locality.postal_code, self.ad1_dict['postal_code'])
        self.assertEqual(test.address.locality.state.name, self.ad1_dict['state'])
        self.assertEqual(test.address.locality.state.code, self.ad1_dict['state_code'])
        self.assertEqual(test.address.locality.state.country.name, self.ad1_dict['country'])
        self.assertEqual(test.address.locality.state.country.code, self.ad1_dict['country_code'])

class FormAddressFieldTestCase(TestCase):

    def setUp(self):
        self.form = TestForm()
        self.field = self.form.base_fields['address']

    def test_to_python_none(self):
        self.assertEqual(self.field.to_python(None), None)

    def test_to_python_empty(self):
        self.assertEqual(self.field.to_python(''), None)

    def test_to_python_invalid_lat_lng(self):
        self.assertRaises(ValidationError, self.field.to_python, {'latitude': 'x'})
        self.assertRaises(ValidationError, self.field.to_python, {'longitude': 'x'})

    def test_to_python(self):
        res = self.field.to_python({'raw': 'Someplace'})
        self.assertEqual(res.raw, 'Someplace')
<|endoftext|>