<commit_before>import enum
import pathlib
import typing  # noqa
import warnings

from pytest import mark, raises

from settei.base import (ConfigKeyError, ConfigTypeError,
                         Configuration, ConfigValueError, ConfigWarning,
                         config_object_property, config_property,
                         get_union_types)

class Enum1(enum.Enum):
    apple = 'apple'
    banana = 'banana'
    cherry = 'cherry'

class Enum2(enum.Enum):
    apple = 'apple'
    bear = 'bear'
    candy = 'candy'

def test_get_union_types():
    assert get_union_types(typing.Union[int, str]) == (int, str)
    assert get_union_types(typing.Union[bool, list, set]) == (bool, list, set)
    assert get_union_types(str) is None
    assert get_union_types('asdf') is None

def test_package_level_api_compatibility():
    import settei
    assert settei.Configuration is Configuration
    assert settei.ConfigWarning is ConfigWarning
    assert settei.config_property is config_property

class TestConfig(dict):
    depth1_required = config_property('key', int)
    depth1_optional = config_property('key', int, default='')
    depth1_warn = config_property('key', int, default='', default_warning=True)
    depth1_default_func = config_property(
        'key', int,
        default_func=lambda self: self.depth1_optional
    )
    depth1_default_func_warn = config_property(
        'key', int,
        default_func=lambda self: self.depth1_optional,
        default_warning=True
    )
    depth2_required = config_property('section.key', str)
    depth2_optional = config_property('section.key', str, default=None)
    depth2_warn = config_property('section.key', str,
                                  default=None, default_warning=True)
    union = config_property('union', typing.Union[int, str])

class EnumTestConfig(dict):
    enum = config_property('enum', Enum1, default=Enum1.apple)
    enum_union = config_property('enum_union', typing.Union[Enum1, Enum2, str])

class TestAppConfig(Configuration):
    database_url = config_property(
        'database.url', str,
        default='sqlite:///test_app.db', default_warning=True
    )

@mark.parametrize('union_value', [123, 'string'])
def test_config_property(union_value: typing.Union[int, str]):
    c = TestConfig(key=123, section={'key': 'val'}, union=union_value)
    assert c.depth1_required == 123
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_optional == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_warn == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func_warn == 123
        assert len(w) == 0
    assert c.depth2_required == 'val'
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional == 'val'
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn == 'val'
        assert len(w) == 0
    assert c.union == union_value

def test_enum_config_property():
    c1 = EnumTestConfig(enum='apple', enum_union='banana')
    c2 = EnumTestConfig(enum='banana', enum_union='candy')
    c3 = EnumTestConfig(enum='dragonfruit', enum_union='apple')
    c4 = EnumTestConfig(enum='cherry', enum_union='foo')
    assert c1.enum == Enum1.apple
    assert c2.enum == Enum1.banana
    assert c1.enum_union == Enum1.banana
    assert c2.enum_union == Enum2.candy
    with raises(ConfigTypeError) as ex:
        c3.enum
    assert ex.value.args[0] == 'Invalid value dragonfruit in <enum \'Enum1\'>. Candidates are: apple, banana, cherry'  # noqa
    with raises(ConfigTypeError) as ex:
        c3.enum_union
    assert ex.value.args[0] == 'Ambiguous enum type for value apple: <Enum1.apple: \'apple\'>, <Enum2.apple: \'apple\'>'  # noqa
    assert c4.enum_union == 'foo'

def test_config_property_absence():
    c = TestConfig()
    with raises(ConfigKeyError):
        c.depth1_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_optional == ''
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_warn == ''
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func == ''
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func_warn == ''
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with raises(ConfigKeyError):
        c.depth2_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional is None
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn is None
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with raises(ConfigKeyError):
        c.union

def test_config_property_absence_2nd_depth():
    c = TestConfig(section={})
    with raises(ConfigKeyError):
        c.depth2_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional is None
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn is None
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)

def test_config_property_type_error():
    c = TestConfig(key='not an integer')
    with raises(ConfigTypeError):
        c.depth1_required
    c2 = TestConfig()
    assert isinstance(c2.depth1_optional, str), \
        'default should be possible to bypass typecheck'
    assert isinstance(c2.depth1_default_func, str), \
        'default_func should be possible to bypass typecheck'

class SampleInterface:
    pass

class Impl(SampleInterface):

    def __init__(self, *args, **kwargs) -> None:
        self.args = args
        self.kwargs = kwargs

    def __repr__(self) -> str:
        args = ', '.join(repr(arg) for arg in self.args)
        kwargs = ', '.join('{0}={1!r}'.format(k, v)
                           for k, v in self.kwargs.items())
        if kwargs:
            args += (', ' if args else '') + kwargs
        return '{0.__module__}.{0.__name__}({1})'.format(type(self), args)

class TestAppConfigObject(Configuration):
    no_default = config_object_property('sample.a', SampleInterface)
    with_default = config_object_property('sample.b', SampleInterface,
                                          default=Impl('default'))
    recursive = config_object_property('sample.c', SampleInterface,
                                       recurse=True)
    cached = config_object_property('sample.a', SampleInterface, cached=True)

def test_config_object_property():
    c = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': ['a', 'b', 'c'],
            'd': 4,
            'e': 5,
            'f': {
                'class': __name__ + ':Impl',
                'nested': True,
                'recursive': {
                    'class': __name__ + ':Impl',
                    'nested': True,
                }
            },
        },
    })
    v = c.no_default
    assert isinstance(v, Impl)
    assert v.args == ('a', 'b', 'c')
    assert v.kwargs == {
        'd': 4, 'e': 5,
        'f': {
            'class': __name__ + ':Impl',
            'nested': True,
            'recursive': {
                'class': __name__ + ':Impl',
                'nested': True,
            }
        },
    }
    assert c.no_default is not c.no_default

def test_config_object_property_recurse():
    c = TestAppConfigObject(sample={
        'a': {'class': __name__ + ':Impl'},
        'c': {
            'class': __name__ + ':Impl',
            '*': ['a', 'b', 'c'],
            'd': 4,
            'e': 5,
            'f': {
                'class': __name__ + ':Impl',
                'nested': True,
                'recursive': {
                    'class': __name__ + ':Impl',
                    'nested': True,
                }
            },
        },
    })
    v = c.recursive
    assert isinstance(v, Impl)
    assert v.args == ('a', 'b', 'c')
    assert frozenset(v.kwargs) == frozenset({'d', 'e', 'f'})
    assert v.kwargs['d'] == 4
    assert v.kwargs['e'] == 5
    f = v.kwargs['f']
    assert isinstance(f, Impl)
    assert f.args == ()
    assert frozenset(f.kwargs) == frozenset({'nested', 'recursive'})
    assert f.kwargs['nested'] is True
    r = f.kwargs['recursive']
    assert isinstance(r, Impl)
    assert r.args == ()
    assert r.kwargs == {'nested': True}

def test_config_object_property_not_dict_error():
    c = TestAppConfigObject(sample={'a': 'Not dict'})
    with raises(ConfigTypeError):
        c.no_default

def test_config_object_property_no_class_field_error():
    c = TestAppConfigObject(sample={'a': {}})
    with raises(ConfigValueError):
        c.no_default

def test_config_object_property_invalid_class_field_error():
    invalid_syntax = TestAppConfigObject(sample={
        'a': {'class': 'invalid.import.path'},
    })
    with raises(ConfigValueError):
        invalid_syntax.no_default
    invalid_type = TestAppConfigObject(sample={
        'a': {'class': 'os.path:supports_unicode_filenames'},
    })
    with raises(ConfigValueError):
        invalid_type.no_default

def test_config_object_property_invalid_args_error():
    non_list = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': 123,  # must be a list, but it's an integer, which is wrong
        },
    })
    with raises(ConfigValueError):
        non_list.no_default
    string = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': 'abc',  # string is not valid either
        },
    })
    with raises(ConfigValueError):
        string.no_default

def test_app_from_file(tmpdir):
    path = tmpdir.join('cfg.toml')
    path.write('''
    [database]
    url = "sqlite:///a.db"
    ''')
    with path.open() as f:
        cfg = TestAppConfig.from_file(f)
    assert cfg.database_url == 'sqlite:///a.db'

def test_app_from_path(tmpdir):
    path = tmpdir.join('cfg.toml')
    path.write('''
    [database]
    url = "sqlite:///b.db"
    ''')
    cfg = TestAppConfig.from_path(pathlib.Path(path.strpath))
    assert cfg.database_url == 'sqlite:///b.db'

def test_config_object_property_cached():
    c = TestAppConfigObject(sample={
        'a': {'class': __name__ + ':Impl'},
    })
    assert c.cached is c.cached
<commit_msg>Add regression test<commit_after>import enum
import pathlib
import typing  # noqa
import warnings

from pytest import mark, raises

from settei.base import (ConfigKeyError, ConfigTypeError,
                         Configuration, ConfigValueError, ConfigWarning,
                         config_object_property, config_property,
                         get_union_types)

class Enum1(enum.Enum):
    apple = 'apple'
    banana = 'banana'
    cherry = 'cherry'

class Enum2(enum.Enum):
    apple = 'apple'
    bear = 'bear'
    candy = 'candy'

def test_get_union_types():
    assert get_union_types(typing.Union[int, str]) == (int, str)
    assert get_union_types(typing.Union[bool, list, set]) == (bool, list, set)
    assert get_union_types(str) is None
    assert get_union_types('asdf') is None

def test_package_level_api_compatibility():
    import settei
    assert settei.Configuration is Configuration
    assert settei.ConfigWarning is ConfigWarning
    assert settei.config_property is config_property

class TestConfig(dict):
    depth1_required = config_property('key', int)
    depth1_optional = config_property('key', int, default='')
    depth1_warn = config_property('key', int, default='', default_warning=True)
    depth1_default_func = config_property(
        'key', int,
        default_func=lambda self: self.depth1_optional
    )
    depth1_default_func_warn = config_property(
        'key', int,
        default_func=lambda self: self.depth1_optional,
        default_warning=True
    )
    depth2_required = config_property('section.key', str)
    depth2_optional = config_property('section.key', str, default=None)
    depth2_warn = config_property('section.key', str,
                                  default=None, default_warning=True)
    union = config_property('union', typing.Union[int, str])

class EnumTestConfig(dict):
    enum = config_property('enum', Enum1, default=Enum1.apple)
    enum_union = config_property('enum_union', typing.Union[Enum1, Enum2, str])

class TestAppConfig(Configuration):
    database_url = config_property(
        'database.url', str,
        default='sqlite:///test_app.db', default_warning=True
    )

@mark.parametrize('union_value', [123, 'string'])
def test_config_property(union_value: typing.Union[int, str]):
    c = TestConfig(key=123, section={'key': 'val'}, union=union_value)
    assert c.depth1_required == 123
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_optional == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_warn == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func == 123
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func_warn == 123
        assert len(w) == 0
    assert c.depth2_required == 'val'
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional == 'val'
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn == 'val'
        assert len(w) == 0
    assert c.union == union_value

def test_enum_config_property():
    c1 = EnumTestConfig(enum='apple', enum_union='banana')
    c2 = EnumTestConfig(enum='banana', enum_union='candy')
    c3 = EnumTestConfig(enum='dragonfruit', enum_union='apple')
    c4 = EnumTestConfig(enum='cherry', enum_union='foo')
    assert c1.enum == Enum1.apple
    assert c2.enum == Enum1.banana
    assert c1.enum_union == Enum1.banana
    assert c2.enum_union == Enum2.candy
    with raises(ConfigTypeError) as ex:
        c3.enum
    assert ex.value.args[0] == 'Invalid value dragonfruit in <enum \'Enum1\'>. Candidates are: apple, banana, cherry'  # noqa
    with raises(ConfigTypeError) as ex:
        c3.enum_union
    assert ex.value.args[0] == 'Ambiguous enum type for value apple: <Enum1.apple: \'apple\'>, <Enum2.apple: \'apple\'>'  # noqa
    assert c4.enum_union == 'foo'

def test_config_property_absence():
    c = TestConfig()
    with raises(ConfigKeyError):
        c.depth1_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_optional == ''
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_warn == ''
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func == ''
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth1_default_func_warn == ''
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with raises(ConfigKeyError):
        c.depth2_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional is None
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn is None
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)
    with raises(ConfigKeyError):
        c.union

def test_config_property_absence_2nd_depth():
    c = TestConfig(section={})
    with raises(ConfigKeyError):
        c.depth2_required
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_optional is None
        assert len(w) == 0
    with warnings.catch_warnings(record=True) as w:
        assert c.depth2_warn is None
        assert len(w) == 1
        assert issubclass(w[-1].category, ConfigWarning)

def test_config_property_type_error():
    c = TestConfig(key='not an integer')
    with raises(ConfigTypeError):
        c.depth1_required
    c2 = TestConfig()
    assert isinstance(c2.depth1_optional, str), \
        'default should be possible to bypass typecheck'
    assert isinstance(c2.depth1_default_func, str), \
        'default_func should be possible to bypass typecheck'

class SampleInterface:
    pass

class Impl(SampleInterface):

    def __init__(self, *args, **kwargs) -> None:
        self.args = args
        self.kwargs = kwargs

    def __repr__(self) -> str:
        args = ', '.join(repr(arg) for arg in self.args)
        kwargs = ', '.join('{0}={1!r}'.format(k, v)
                           for k, v in self.kwargs.items())
        if kwargs:
            args += (', ' if args else '') + kwargs
        return '{0.__module__}.{0.__name__}({1})'.format(type(self), args)

class TestAppConfigObject(Configuration):
    no_default = config_object_property('sample.a', SampleInterface)
    with_default = config_object_property('sample.b', SampleInterface,
                                          default=Impl('default'))
    recursive = config_object_property('sample.c', SampleInterface,
                                       recurse=True)
    cached = config_object_property('sample.a', SampleInterface, cached=True)
    cached_with_default = config_object_property('sample.b', SampleInterface,
                                                 default=Impl('default'),
                                                 cached=True)

def test_config_object_property():
    c = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': ['a', 'b', 'c'],
            'd': 4,
            'e': 5,
            'f': {
                'class': __name__ + ':Impl',
                'nested': True,
                'recursive': {
                    'class': __name__ + ':Impl',
                    'nested': True,
                }
            },
        },
    })
    v = c.no_default
    assert isinstance(v, Impl)
    assert v.args == ('a', 'b', 'c')
    assert v.kwargs == {
        'd': 4, 'e': 5,
        'f': {
            'class': __name__ + ':Impl',
            'nested': True,
            'recursive': {
                'class': __name__ + ':Impl',
                'nested': True,
            }
        },
    }
    assert c.no_default is not c.no_default
    vv = c.with_default
    assert isinstance(vv, Impl)
    assert vv.args == ('default',)
    assert vv.kwargs == {}

def test_config_object_property_recurse():
    c = TestAppConfigObject(sample={
        'a': {'class': __name__ + ':Impl'},
        'c': {
            'class': __name__ + ':Impl',
            '*': ['a', 'b', 'c'],
            'd': 4,
            'e': 5,
            'f': {
                'class': __name__ + ':Impl',
                'nested': True,
                'recursive': {
                    'class': __name__ + ':Impl',
                    'nested': True,
                }
            },
        },
    })
    v = c.recursive
    assert isinstance(v, Impl)
    assert v.args == ('a', 'b', 'c')
    assert frozenset(v.kwargs) == frozenset({'d', 'e', 'f'})
    assert v.kwargs['d'] == 4
    assert v.kwargs['e'] == 5
    f = v.kwargs['f']
    assert isinstance(f, Impl)
    assert f.args == ()
    assert frozenset(f.kwargs) == frozenset({'nested', 'recursive'})
    assert f.kwargs['nested'] is True
    r = f.kwargs['recursive']
    assert isinstance(r, Impl)
    assert r.args == ()
    assert r.kwargs == {'nested': True}

def test_config_object_property_not_dict_error():
    c = TestAppConfigObject(sample={'a': 'Not dict'})
    with raises(ConfigTypeError):
        c.no_default

def test_config_object_property_no_class_field_error():
    c = TestAppConfigObject(sample={'a': {}})
    with raises(ConfigValueError):
        c.no_default

def test_config_object_property_invalid_class_field_error():
    invalid_syntax = TestAppConfigObject(sample={
        'a': {'class': 'invalid.import.path'},
    })
    with raises(ConfigValueError):
        invalid_syntax.no_default
    invalid_type = TestAppConfigObject(sample={
        'a': {'class': 'os.path:supports_unicode_filenames'},
    })
    with raises(ConfigValueError):
        invalid_type.no_default

def test_config_object_property_invalid_args_error():
    non_list = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': 123,  # must be a list, but it's an integer, which is wrong
        },
    })
    with raises(ConfigValueError):
        non_list.no_default
    string = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': 'abc',  # string is not valid either
        },
    })
    with raises(ConfigValueError):
        string.no_default

def test_app_from_file(tmpdir):
    path = tmpdir.join('cfg.toml')
    path.write('''
    [database]
    url = "sqlite:///a.db"
    ''')
    with path.open() as f:
        cfg = TestAppConfig.from_file(f)
    assert cfg.database_url == 'sqlite:///a.db'

def test_app_from_path(tmpdir):
    path = tmpdir.join('cfg.toml')
    path.write('''
    [database]
    url = "sqlite:///b.db"
    ''')
    cfg = TestAppConfig.from_path(pathlib.Path(path.strpath))
    assert cfg.database_url == 'sqlite:///b.db'

def test_config_object_property_cached():
    c = TestAppConfigObject(sample={
        'a': {
            'class': __name__ + ':Impl',
            '*': ['a', 'b', 'c'],
            'd': 4,
            'e': 5,
        },
    })
    assert c.cached is c.cached
    assert isinstance(c.cached, Impl)
    assert c.cached.args == ('a', 'b', 'c')
    assert c.cached.kwargs == {'d': 4, 'e': 5}
    assert isinstance(c.cached_with_default, Impl)
    assert c.cached_with_default.args == ('default',)
    assert c.cached_with_default.kwargs == {}
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
This module builds HTML documentation for models by introspecting the model
unit tests.
"""

from django.core import meta, template
import runtests
import inspect, os, re, sys

MODEL_DOC_TEMPLATE = """

{{ title }}

{{ blurb }}

Model source code

{{ model_source }}

Sample API usage

{{ api_usage }}

"""

def make_docs_from_model_tests(output_dir):
    from django.conf import settings

    # Manually set INSTALLED_APPS to point to the test app.
    settings.INSTALLED_APPS = (runtests.APP_NAME,)

    for model_name in runtests.get_test_models():
        mod = meta.get_app(model_name)

        # Clean up the title and blurb.
        title, blurb = mod.__doc__.strip().split('\n', 1)
        blurb = '
%s
' % blurb.strip().replace('\n\n', '

')
        api_usage = mod.API_TESTS

        # Get the source code of the model, without the docstring or the
        # API_TESTS variable.
        model_source = inspect.getsource(mod)
        model_source = model_source.replace(mod.__doc__, '')
        model_source = model_source.replace(mod.API_TESTS, '')
        model_source = model_source.replace('""""""\n', '\n')
        model_source = model_source.replace('API_TESTS = ', '')
        model_source = model_source.strip()

        # Run this through the template system.
        t = template.Template(MODEL_DOC_TEMPLATE)
        c = template.Context(locals())
        html = t.render(c)

        file_name = os.path.join(output_dir, 'model_' + model_name + '.html')
        try:
            fp = open(file_name, 'w')
        except IOError:
            sys.stderr.write("Couldn't write to %s.\n" % file_name)
            continue
        else:
            fp.write(html)
            fp.close()

if __name__ == "__main__":
    import sys
    make_docs_from_model_tests(sys.argv[1])
<commit_msg>Changed template in tests.builddocs<commit_after>#!/usr/bin/env python
"""
This module builds HTML documentation for models by introspecting the model
unit tests.
"""

from django.core import meta, template
import runtests
import inspect, os, re, sys

MODEL_DOC_TEMPLATE = """

{{ title }}

{{ blurb }}

Model source code

{{ model_source }}

Sample API usage

{{ api_usage }}

"""

def make_docs_from_model_tests(output_dir):
    from django.conf import settings

    # Manually set INSTALLED_APPS to point to the test app.
    settings.INSTALLED_APPS = (runtests.APP_NAME,)

    for model_name in runtests.get_test_models():
        mod = meta.get_app(model_name)

        # Clean up the title and blurb.
        title, blurb = mod.__doc__.strip().split('\n', 1)
        blurb = '
%s
' % blurb.strip().replace('\n\n', '

')
        api_usage = mod.API_TESTS

        # Get the source code of the model, without the docstring or the
        # API_TESTS variable.
        model_source = inspect.getsource(mod)
        model_source = model_source.replace(mod.__doc__, '')
        model_source = model_source.replace(mod.API_TESTS, '')
        model_source = model_source.replace('""""""\n', '\n')
        model_source = model_source.replace('API_TESTS = ', '')
        model_source = model_source.strip()

        # Run this through the template system.
        t = template.Template(MODEL_DOC_TEMPLATE)
        c = template.Context(locals())
        html = t.render(c)

        file_name = os.path.join(output_dir, 'model_' + model_name + '.html')
        try:
            fp = open(file_name, 'w')
        except IOError:
            sys.stderr.write("Couldn't write to %s.\n" % file_name)
            continue
        else:
            fp.write(html)
            fp.close()

if __name__ == "__main__":
    import sys
    make_docs_from_model_tests(sys.argv[1])
<|endoftext|>


<commit_before>import re, os
from django.template import Library
from sorl.thumbnail import Thumbnail, METHOD_LIST
from django.conf import settings

register = Library()

@register.filter
def thumbnail_url(filename, arg=""):
    try:
        thumbnail = get_thumbnail(filename, arg)
    except:
        if hasattr(settings, 'THUMBNAIL_DEBUG') and settings.THUMBNAIL_DEBUG:
            raise
        else:
            return ""

    return thumbnail.get_url()

def get_thumbnail(filename, arg=""):

    kwargs  = {\
        'prefix': '',
        'subdir': '_thumbs',
        'quality': 85,
    }

    for a in kwargs.keys():
        if hasattr(settings, 'THUMBNAIL_%s' % a.upper()):
            kwargs.update({a: getattr(settings, 'THUMBNAIL_%s' % a.upper())})

    kwargs.update({\
        'filename': filename,
        'size': (80,80),
        'crop': False,
        'upscale': False,
        'bw': False,
    })

    for m in METHOD_LIST:
        if arg.find(m) != -1:
            kwargs.update({m: True})
    quality_pat = re.compile(r'q(\d+)')
    m = quality_pat.search(arg)
    if m:
        quality = int(m.group(1))
        if quality > 0 and quality <= 100:
            kwargs.update(quality=quality)
    size_pat = re.compile(r'(\d+)x(\d+)')
    m = size_pat.search(arg)
    if m:
        kwargs.update(size=(int(m.group(1)), int(m.group(2))))

    return Thumbnail(**kwargs)
<commit_msg>...<commit_after>import re, os
from django.template import Library
from sorl.thumbnail import Thumbnail, METHOD_LIST
from django.conf import settings

register = Library()

@register.filter
def thumbnail_url(filename, arg=""):
    try:
        thumbnail = get_thumbnail(filename, arg)
    except:
        if hasattr(settings, 'THUMBNAIL_DEBUG') and settings.THUMBNAIL_DEBUG:
            raise
        else:
            return ""

    return thumbnail.get_url()

def get_thumbnail(filename, arg=""):

    kwargs  = {\
        'prefix': '',
        'subdir': '_thumbs',
        'quality': 85,
    }

    for a in kwargs.keys():
        if hasattr(settings, 'THUMBNAIL_%s' % a.upper()):
            kwargs.update({a: getattr(settings, 'THUMBNAIL_%s' % a.upper())})

    kwargs.update({'filename': filename, 'size': (80,80)})

    for m in METHOD_LIST:
        if arg.find(m) != -1:
            kwargs.update({m: True})
        else:
            kwargs.update({m: False})
    quality_pat = re.compile(r'q(\d+)')
    m = quality_pat.search(arg)
    if m:
        quality = int(m.group(1))
        if quality > 0 and quality <= 100:
            kwargs.update(quality=quality)
    size_pat = re.compile(r'(\d+)x(\d+)')
    m = size_pat.search(arg)
    if m:
        kwargs.update(size=(int(m.group(1)), int(m.group(2))))

    return Thumbnail(**kwargs)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
This file contains large scale bulk tests to check whether the code is
running find or not. This is tested using a large number of files from
commons.wikimedia.org.
"""

from __future__ import (division, absolute_import, unicode_literals,
                        print_function)

import datetime
import json
import os
import ssl
import subprocess
import types
from collections import Counter

import pytest
from retry import retry
from six import string_types
from six.moves.urllib.error import URLError

from file_metadata._compat import makedirs
from file_metadata.generic_file import GenericFile
from file_metadata.image.image_file import ImageFile
from file_metadata.utilities import download
from tests import is_travis, unittest, CACHE_DIR

try:
    import pywikibot
except ImportError as err:
    raise unittest.SkipTest("Module pywikibot not found. Please install it "
                            "to run bulk tests.")
except RuntimeError as err:
    pywikibot = err
    raise unittest.SkipTest('pywikibot requires a user-config.py, which was '
                            'not found. Please create one and then run '
                            'bulk tests.')

from pywikibot import pagegenerators

def dump_log(data, logname, _type='text', header=None):
    """
    Dump the given info into a mediawiki file in the User namespace
    'User:<Username>/logs/<logname>'.

    :param data:    The data to log.
    :param logname: The name of the page to save it in.
    :param _type:   The type of data being logged. json or text.
    :param header:  The header to add to the page before dumping data.
    """
    if data is not None:
        if _type == 'json':
            logs = '
{0}
'.format(json.dumps(
                data, sort_keys=True, indent=2, separators=(',', ': ')))
        elif _type == 'text':
            if isinstance(data, string_types):
                logs = data
            elif isinstance(data, (tuple, list, types.GeneratorType)):
                logs = "\n".join(data)
            else:
                raise ValueError('Unexpected output got. Expected str, tuple, '
                                 'list or generator when _type=text is given.')
        else:
            raise ValueError('Unexpected value "{0} given."'.format(_type))

        if header is not None:
            logs = header + '\n' + logs

        site = pywikibot.Site()
        site.login()

        page = pywikibot.Page(
            site, 'User:' + site.username() + '/logs/' + logname)
        page.put(logs, "Logged with file-metadata's dump_log()")

@pytest.mark.timeout(60 * 60)
class PyWikiBotTestHelper(unittest.TestCase):

    def setUp(self):
        self.site = pywikibot.Site()
        self.site.login()

        makedirs(CACHE_DIR, exist_ok=True)

    @staticmethod
    @retry((ssl.SSLError, URLError), tries=5)
    def download_page(page, fname=None):
        url = page.fileUrl()
        fname = fname or page.title(as_filename=True).encode('ascii',
                                                             'replace')
        fpath = os.path.join(CACHE_DIR, fname)
        download(url, fpath)
        return fpath

    def factory(self, args, fname=None):
        """
        Use pywikibot to fetch pages based on the given arguments. If the
        system is detected to be TRAVIS, it auto deletes the file afterwards.

        :param args:  The args to give to pywikibot
        :param fname: If this is given, the file is stored in this filename.
        :return:      A generator with the pages asked
        """
        gen_factory = pagegenerators.GeneratorFactory(self.site)
        for arg in args:
            gen_factory.handleArg(arg)

        generator = gen_factory.getCombinedGenerator()
        if not generator:
            self.fail('No generator was asked from the factory.')
        else:
            pregen = pagegenerators.PreloadingGenerator(generator,
                                                        groupsize=50)
            for page in pregen:
                if page.exists() and not page.isRedirectPage():
                    page_path = self.download_page(page, fname=fname)
                    yield page, page_path
                    if is_travis():
                        os.remove(page_path)

class BulkCategoryTest(PyWikiBotTestHelper):

    def _test_file(self, page, path):
        log = []
        stats = {'mime': None}
        txt, im = [], []
        # First cell has text info
        # Second cell has image bounding boxes if applicable

        _file = GenericFile.create(path)
        log.append('===== {0} ====='.format(
                   page.title(asLink=True, textlink=True)))

        mime = _file.analyze_mimetype().get('File:MIMEType', "ERROR")
        stats['mime'] = mime
        txt.append("* '''Mime Type''': " + mime)
        if mime == 'application/ogg':
            _type = _file.analyze_exiftool().get('File:FileType', 'ERROR')
            txt.append("* '''File type''': " + _type)

        if isinstance(_file, ImageFile) and mime in {
                'image/jpeg', 'image/png', 'image/gif', 'image/tiff'}:
            # Can read the pixel data of the image ...

            height, width = _file.fetch('ndarray').shape[:2]
            if width > height:
                scale = min(200, width) * 1.0 / width
            else:
                scale = min(200, height) * 1.0 / height
            scaled_h, scaled_w = int(height * scale), int(width * scale)

            im.append('
')
            im.append('{0}|{1}x{2}px]]'.format(
                      page.title(underscore=False, asLink=True)[:-2],
                      scaled_w, scaled_h))
            box = ('

')

            # Analyze barcodes
            barcode_data = {}
            try:
                barcode_data = _file.analyze_barcode()
            except subprocess.CalledProcessError:
                txt.append("* '''Barcodes ERROR:''' Unsupported file type")
            barcodes = barcode_data.get('zxing:Barcodes', None)
            if barcodes is not None and len(barcodes) > 0:
                txt.append("* '''Barcodes found:'''")
                for bar in barcodes:
                    txt.append("** Data: " + bar['data'])
                    txt.append("** Format: " + str(bar['format']))
                    txt.append("** Position: " + str(bar['points']))
                    if len(bar['points']) == 2:
                        # left, right
                        l, r = [(int(i * scale), int(j * scale))
                                for (i, j) in bar['points']]
                        im.append(box.format(
                            cls='barcode', col='ff0000', l=l[0], t=l[1],
                            w=r[0] - l[0] + 2, h=r[1] - l[1] + 2))
                    elif len(bar['points']) == 4:
                        # bottomLeft, topLeft, topRight, bottomRight
                        lb, lt, rt, rb = [(int(i * scale), int(j * scale))
                                          for (i, j) in bar['points']]
                        im.append(box.format(
                            cls='barcode', col='ff0000',
                            l=min(lb[0], lt[0]), t=min(lt[1], rt[1]),
                            w=max(rb[0] - lb[0], rt[0] - lt[0]),
                            h=max(rb[1] - rt[1], lb[1] - lt[1])))

            # Make a table for text and image
            log += ['{| class="wikitable"', '|'] + txt + ['|'] + im + ['|}']
        else:
            # Image cell is empty as it's not an image
            log += ['{| class="wikitable"', '|'] + txt + ['|', '', '|}']
        return log, stats

    def _test_category(self, category, limit=10000):
        log = []
        cat = pywikibot.Category(self.site, category)

        stats = {'count': 0, 'mime': [], 'ext': [],
                 'start_time': datetime.datetime.now()}
        for count, (page, path) in enumerate(self.factory(
                ['-catr:' + cat.title(withNamespace=False, underscore=True),
                 '-limit:' + str(limit),
                 '-ns:File'])):

            print(count + 1, '. Analyzing', page.title(underscore=False))
            if (count + 1) % 500 == 1:
                log.append('== {0} to {1} =='.format(count + 1, count + 500))
            _log, _stat = self._test_file(page, path)
            log += _log

            stats['count'] = count + 1
            stats['ext'].append(os.path.splitext(page.title())[-1])
            stats['mime'].append(_stat['mime'])

        stats['end_time'] = datetime.datetime.now()
        stats['timetaken'] = (stats['end_time'] -
                              stats['start_time']).total_seconds()

        title = 'Category {0}'.format(cat.title(underscore=False,
                                                withNamespace=False))
        summary = ["Analyis from files of " + cat.title(asLink=True,
                                                        textlink=True),
                   "* '''Time taken''': " + str(stats['timetaken']),
                   "* '''Files analyzed''': " + str(stats['count']),
                   "* '''MIME type stats''': "]
        for mime, count in Counter(stats['mime']).items():
            summary.append("** " + str(mime) + " - " + str(count))
        summary.append("* '''Enxtension Stats''': ")
        for ext, count in Counter(stats['ext']).items():
            summary.append("** " + str(ext) + " - " + str(count))

        dump_log(log, logname=title, header="\n".join(summary))

    def test_png_files(self):
        self._test_category('PNG files')

    def test_svg_files(self):
        self._test_category('SVG files')

    def test_jpeg_files(self):
        self._test_category('JPEG files')

    def test_gif_files(self):
        self._test_category('GIF files')

    def test_tiff_files(self):
        self._test_category('TIFF files')

    def test_ogv_videos(self):
        self._test_category('Ogv videos')

    def test_animated_gif_files(self):
        self._test_category('Animated GIF files')

    def test_animated_png(self):
        self._test_category('Animated PNG')

    def test_animated_svg(self):
        self._test_category('Animated SVG')

    def test_flac_files(self):
        self._test_category('FLAC files')

    def test_wav_files(self):
        self._test_category('WAV files')

    def test_ogg_sound_files(self):
        self._test_category('Ogg sound files')

    def test_midi_files(self):
        self._test_category('MIDI files')

    def test_djvu_files(self):
        self._test_category('DjVu files')

    def test_xcf_files(self):
        self._test_category('XCF files')

    def test_pdf_files(self):
        self._test_category('PDF files')

    def test_barcode(self):
        self._test_category('Barcode')
<commit_msg>bulk_test: Timeout is download takes >15 min<commit_after># -*- coding: utf-8 -*-
"""
This file contains large scale bulk tests to check whether the code is
running find or not. This is tested using a large number of files from
commons.wikimedia.org.
"""

from __future__ import (division, absolute_import, unicode_literals,
                        print_function)

import datetime
import json
import os
import ssl
import subprocess
import types
from collections import Counter

import pytest
from retry import retry
from six import string_types
from six.moves.urllib.error import URLError

from file_metadata._compat import makedirs
from file_metadata.generic_file import GenericFile
from file_metadata.image.image_file import ImageFile
from file_metadata.utilities import download
from tests import is_travis, unittest, CACHE_DIR

try:
    import pywikibot
except ImportError as err:
    raise unittest.SkipTest("Module pywikibot not found. Please install it "
                            "to run bulk tests.")
except RuntimeError as err:
    pywikibot = err
    raise unittest.SkipTest('pywikibot requires a user-config.py, which was '
                            'not found. Please create one and then run '
                            'bulk tests.')

from pywikibot import pagegenerators

def dump_log(data, logname, _type='text', header=None):
    """
    Dump the given info into a mediawiki file in the User namespace
    'User:<Username>/logs/<logname>'.

    :param data:    The data to log.
    :param logname: The name of the page to save it in.
    :param _type:   The type of data being logged. json or text.
    :param header:  The header to add to the page before dumping data.
    """
    if data is not None:
        if _type == 'json':
            logs = '
{0}
'.format(json.dumps(
                data, sort_keys=True, indent=2, separators=(',', ': ')))
        elif _type == 'text':
            if isinstance(data, string_types):
                logs = data
            elif isinstance(data, (tuple, list, types.GeneratorType)):
                logs = "\n".join(data)
            else:
                raise ValueError('Unexpected output got. Expected str, tuple, '
                                 'list or generator when _type=text is given.')
        else:
            raise ValueError('Unexpected value "{0} given."'.format(_type))

        if header is not None:
            logs = header + '\n' + logs

        site = pywikibot.Site()
        site.login()

        page = pywikibot.Page(
            site, 'User:' + site.username() + '/logs/' + logname)
        page.put(logs, "Logged with file-metadata's dump_log()")

@pytest.mark.timeout(60 * 60)
class PyWikiBotTestHelper(unittest.TestCase):

    def setUp(self):
        self.site = pywikibot.Site()
        self.site.login()

        makedirs(CACHE_DIR, exist_ok=True)

    @staticmethod
    @retry((ssl.SSLError, URLError), tries=5)
    def download_page(page, fname=None, timeout=None):
        url = page.fileUrl()
        fname = fname or page.title(as_filename=True).encode('ascii',
                                                             'replace')
        fpath = os.path.join(CACHE_DIR, fname)
        download(url, fpath, timeout=timeout)
        return fpath

    def factory(self, args, fname=None):
        """
        Use pywikibot to fetch pages based on the given arguments. If the
        system is detected to be TRAVIS, it auto deletes the file afterwards.

        :param args:  The args to give to pywikibot
        :param fname: If this is given, the file is stored in this filename.
        :return:      A generator with the pages asked
        """
        gen_factory = pagegenerators.GeneratorFactory(self.site)
        for arg in args:
            gen_factory.handleArg(arg)

        generator = gen_factory.getCombinedGenerator()
        if not generator:
            self.fail('No generator was asked from the factory.')
        else:
            pregen = pagegenerators.PreloadingGenerator(generator,
                                                        groupsize=50)
            for page in pregen:
                if page.exists() and not page.isRedirectPage():
                    try:
                        page_path = self.download_page(page, fname=fname,
                                                       timeout=15 * 60)
                    except URLError:
                        # URLError is raised if the download timesout.
                        continue
                    yield page, page_path
                    if is_travis():
                        os.remove(page_path)

class BulkCategoryTest(PyWikiBotTestHelper):

    def _test_file(self, page, path):
        log = []
        stats = {'mime': None}
        txt, im = [], []
        # First cell has text info
        # Second cell has image bounding boxes if applicable

        _file = GenericFile.create(path)
        log.append('===== {0} ====='.format(
                   page.title(asLink=True, textlink=True)))

        mime = _file.analyze_mimetype().get('File:MIMEType', "ERROR")
        stats['mime'] = mime
        txt.append("* '''Mime Type''': " + mime)
        if mime == 'application/ogg':
            _type = _file.analyze_exiftool().get('File:FileType', 'ERROR')
            txt.append("* '''File type''': " + _type)

        if isinstance(_file, ImageFile) and mime in {
                'image/jpeg', 'image/png', 'image/gif', 'image/tiff'}:
            # Can read the pixel data of the image ...

            height, width = _file.fetch('ndarray').shape[:2]
            if width > height:
                scale = min(200, width) * 1.0 / width
            else:
                scale = min(200, height) * 1.0 / height
            scaled_h, scaled_w = int(height * scale), int(width * scale)

            im.append('
')
            im.append('{0}|{1}x{2}px]]'.format(
                      page.title(underscore=False, asLink=True)[:-2],
                      scaled_w, scaled_h))
            box = ('

')

            # Analyze barcodes
            barcode_data = {}
            try:
                barcode_data = _file.analyze_barcode()
            except subprocess.CalledProcessError:
                txt.append("* '''Barcodes ERROR:''' Unsupported file type")
            barcodes = barcode_data.get('zxing:Barcodes', None)
            if barcodes is not None and len(barcodes) > 0:
                txt.append("* '''Barcodes found:'''")
                for bar in barcodes:
                    txt.append("** Data: " + bar['data'])
                    txt.append("** Format: " + str(bar['format']))
                    txt.append("** Position: " + str(bar['points']))
                    if len(bar['points']) == 2:
                        # left, right
                        l, r = [(int(i * scale), int(j * scale))
                                for (i, j) in bar['points']]
                        im.append(box.format(
                            cls='barcode', col='ff0000', l=l[0], t=l[1],
                            w=r[0] - l[0] + 2, h=r[1] - l[1] + 2))
                    elif len(bar['points']) == 4:
                        # bottomLeft, topLeft, topRight, bottomRight
                        lb, lt, rt, rb = [(int(i * scale), int(j * scale))
                                          for (i, j) in bar['points']]
                        im.append(box.format(
                            cls='barcode', col='ff0000',
                            l=min(lb[0], lt[0]), t=min(lt[1], rt[1]),
                            w=max(rb[0] - lb[0], rt[0] - lt[0]),
                            h=max(rb[1] - rt[1], lb[1] - lt[1])))

            # Make a table for text and image
            log += ['{| class="wikitable"', '|'] + txt + ['|'] + im + ['|}']
        else:
            # Image cell is empty as it's not an image
            log += ['{| class="wikitable"', '|'] + txt + ['|', '', '|}']
        return log, stats

    def _test_category(self, category, limit=10000):
        log = []
        cat = pywikibot.Category(self.site, category)

        stats = {'count': 0, 'mime': [], 'ext': [],
                 'start_time': datetime.datetime.now()}
        for count, (page, path) in enumerate(self.factory(
                ['-catr:' + cat.title(withNamespace=False, underscore=True),
                 '-limit:' + str(limit),
                 '-ns:File'])):

            print(count + 1, '. Analyzing', page.title(underscore=False))
            if (count + 1) % 500 == 1:
                log.append('== {0} to {1} =='.format(count + 1, count + 500))
            _log, _stat = self._test_file(page, path)
            log += _log

            stats['count'] = count + 1
            stats['ext'].append(os.path.splitext(page.title())[-1])
            stats['mime'].append(_stat['mime'])

        stats['end_time'] = datetime.datetime.now()
        stats['timetaken'] = (stats['end_time'] -
                              stats['start_time']).total_seconds()

        title = 'Category {0}'.format(cat.title(underscore=False,
                                                withNamespace=False))
        summary = ["Analyis from files of " + cat.title(asLink=True,
                                                        textlink=True),
                   "* '''Time taken''': " + str(stats['timetaken']),
                   "* '''Files analyzed''': " + str(stats['count']),
                   "* '''MIME type stats''': "]
        for mime, count in Counter(stats['mime']).items():
            summary.append("** " + str(mime) + " - " + str(count))
        summary.append("* '''Enxtension Stats''': ")
        for ext, count in Counter(stats['ext']).items():
            summary.append("** " + str(ext) + " - " + str(count))

        dump_log(log, logname=title, header="\n".join(summary))

    def test_png_files(self):
        self._test_category('PNG files')

    def test_svg_files(self):
        self._test_category('SVG files')

    def test_jpeg_files(self):
        self._test_category('JPEG files')

    def test_gif_files(self):
        self._test_category('GIF files')

    def test_tiff_files(self):
        self._test_category('TIFF files')

    def test_ogv_videos(self):
        self._test_category('Ogv videos')

    def test_animated_gif_files(self):
        self._test_category('Animated GIF files')

    def test_animated_png(self):
        self._test_category('Animated PNG')

    def test_animated_svg(self):
        self._test_category('Animated SVG')

    def test_flac_files(self):
        self._test_category('FLAC files')

    def test_wav_files(self):
        self._test_category('WAV files')

    def test_ogg_sound_files(self):
        self._test_category('Ogg sound files')

    def test_midi_files(self):
        self._test_category('MIDI files')

    def test_djvu_files(self):
        self._test_category('DjVu files')

    def test_xcf_files(self):
        self._test_category('XCF files')

    def test_pdf_files(self):
        self._test_category('PDF files')

    def test_barcode(self):
        self._test_category('Barcode')
<|endoftext|>


<commit_before>src = """
window.pywebview = {
    token: '%s',
    platform: '%s',
    api: {},

    _createApi: function(funcList) {
        for (var i = 0; i < funcList.length; i++) {
            var funcName = funcList[i].func;
            var params = funcList[i].params;

            var funcBody =
                "var id = (Math.random() + '').substring(2); " +
                "var promise = new Promise(function(resolve, reject) { " +
                    "window.pywebview._checkValue('" + funcName + "', resolve, reject, id); " +
                "}); " +
                "window.pywebview._bridge.call('" + funcName + "', JSON.stringify(arguments), id); " +
                "return promise;"

            window.pywebview.api[funcName] = new Function(params, funcBody)
            window.pywebview._returnValues[funcName] = {}
        }
    },

    _bridge: {
        call: function (funcName, params, id) {
            switch(window.pywebview.platform) {
                case 'mshtml':
                case 'cef':
                case 'qtwebkit':
                    return window.external.call(funcName, params, id);
                case 'edgehtml':
                    return window.external.notify(JSON.stringify([funcName, params, id]));
                case 'cocoa':
                    return window.webkit.messageHandlers.jsBridge.postMessage(JSON.stringify([funcName, params, id]));
                case 'qtwebengine':
                    new QWebChannel(qt.webChannelTransport, function(channel) {
                        channel.objects.external.call(funcName, params, id);
                    });
                    break;
                case 'gtk':
                    document.title = JSON.stringify({"type": "invoke", "uid": "%s", "function": funcName, "param": params, "id": id});
                    break;
            }
        }
    },

    _checkValue: function(funcName, resolve, reject, id) {
         var check = setInterval(function () {
            var returnObj = window.pywebview._returnValues[funcName][id];
            if (returnObj) {
                var value = returnObj.value;
                var isError = returnObj.isError;

                delete window.pywebview._returnValues[funcName][id];
                clearInterval(check);

                if (isError) {
                    var pyError = JSON.parse(value);
                    var error = new Error(pyError.message);
                    error.name = pyError.name;
                    error.stack = pyError.stack;

                    reject(error);
                } else {
                    resolve(JSON.parse(value));
                }
            }
         }, 100)
    },

    _returnValues: {}
}
window.pywebview._createApi(%s);
window.dispatchEvent(new CustomEvent('pywebviewready'));
"""
<commit_msg>Fix JS API overwriting function param with the name id<commit_after>src = """
window.pywebview = {
    token: '%s',
    platform: '%s',
    api: {},

    _createApi: function(funcList) {
        for (var i = 0; i < funcList.length; i++) {
            var funcName = funcList[i].func;
            var params = funcList[i].params;

            var funcBody =
                "var __id = (Math.random() + '').substring(2); " +
                "var promise = new Promise(function(resolve, reject) { " +
                    "window.pywebview._checkValue('" + funcName + "', resolve, reject, __id); " +
                "}); " +
                "window.pywebview._bridge.call('" + funcName + "', JSON.stringify(arguments), __id); " +
                "return promise;"

            window.pywebview.api[funcName] = new Function(params, funcBody)
            window.pywebview._returnValues[funcName] = {}
        }
    },

    _bridge: {
        call: function (funcName, params, id) {
            switch(window.pywebview.platform) {
                case 'mshtml':
                case 'cef':
                case 'qtwebkit':
                    return window.external.call(funcName, params, id);
                case 'edgehtml':
                    return window.external.notify(JSON.stringify([funcName, params, id]));
                case 'cocoa':
                    return window.webkit.messageHandlers.jsBridge.postMessage(JSON.stringify([funcName, params, id]));
                case 'qtwebengine':
                    new QWebChannel(qt.webChannelTransport, function(channel) {
                        channel.objects.external.call(funcName, params, id);
                    });
                    break;
                case 'gtk':
                    document.title = JSON.stringify({"type": "invoke", "uid": "%s", "function": funcName, "param": params, "id": id});
                    break;
            }
        }
    },

    _checkValue: function(funcName, resolve, reject, id) {
         var check = setInterval(function () {
            var returnObj = window.pywebview._returnValues[funcName][id];
            if (returnObj) {
                var value = returnObj.value;
                var isError = returnObj.isError;

                delete window.pywebview._returnValues[funcName][id];
                clearInterval(check);

                if (isError) {
                    var pyError = JSON.parse(value);
                    var error = new Error(pyError.message);
                    error.name = pyError.name;
                    error.stack = pyError.stack;

                    reject(error);
                } else {
                    resolve(JSON.parse(value));
                }
            }
         }, 100)
    },

    _returnValues: {}
}
window.pywebview._createApi(%s);
window.dispatchEvent(new CustomEvent('pywebviewready'));
"""
<|endoftext|>


<commit_before>"""
Main solving algorithm of the
Monge-Ampere screen mapping potential problem using
a finite difference scheme.
"""
import numpy as np
import scipy as sp
import scipy.interpolate
import os.path

def setup(ri, li, rs, dxS, image, mask):
    """
    Takes proton radiography distance parameters, a proton flux image, and a
    flux image selection mask and sets up the deflection-field
    potential problem.

    Args:
        ri (float): Distance from source to plasma (cm).
        li (float): Distance across plasma (cm).
        rs (float): Distance from plasma to screen (cm).
        dxS (float): Pixel size (cm)
        image (array): Proton flux image to analyze.
        mask (array): Bitwise mask of data to analyze in the flux image,
                      assuming 1 is masked and 0 is unmasked.

    Returns:
        plasma_x (array): Plasma x-coordinates (cm).
        plasma_y (array): Plasma x-coordinates (cm).
        flux0 (array): Initial flux image.
        flux (array): Flux image.
    """

    magnify = (rs+ri+.5*li)/(ri+.5*li)
    # Leave the original image unchanged.
    flux_image = np.copy(image)
    flux_mask = np.copy(mask)

    # Mask the flux_image and
    # fill the masked positions with the mean flux.
    #
    # TODO This may add gradient to the proton fluence and therefore
    # create artificial magnetic fields-- may need to find a different way to
    # do this. Possibly extrapolate gradients from the boundary?
    #flux = np.ma.array(flux_image, mask=np.logical_not(flux_mask))
    flux = flux_image * np.logical_not(flux_mask)
    flux_mean = np.nanmean(flux)
    flux = flux + flux_mean * flux_mask

    # Calculate the total flux throughout the selected region
    flux_tot = np.nansum(np.logical_not(flux_mask)*flux)

    # TODO
    # @usualgaussnumber applies a Gaussian filter after this step, see
    # pradreconstructionfluxdistsetup.m

    # Prepare the flux and
    # flux_0 for the primary reconstruction process.
    # flux_0 and flux_0 mask are the initial flux distributions,
    # which we choose to be uniform over the selected region while
    # preserving total flux.
    # Then add a shadow flux region to flux0 which is equal to the flux mean.
    flux0_mask = np.copy(flux_mask)
    flux0 = np.zeros(flux.shape)
    flux0 = (flux_tot * np.logical_not(flux0_mask))/(np.nansum(np.logical_not(flux0_mask)))
    flux0 = flux0 + flux0_mask * np.mean(flux0)

    # Variables:
    # Nx, N1 are the number of cells on the 0th and 1st axis
    # dx is the length of each cell (square)
    # L0, L1 are the lengths of the flux image in the 0th and 1st axis
    Nx, Ny = flux.shape
    Lx = dxS*Nx
    Ly = dxS*Ny

    # Create the coordinates for the flux samples for the centers of the
    # bins on the screen.
    screen_x = np.linspace((-Lx/2)+.5*dxS, (Lx/2)-.5*dxS, Nx)
    screen_y = np.linspace((-Ly/2)+.5*dxS, (Ly/2)-.5*dxS, Ny)
    screen_x, screen_y = np.meshgrid(screen_x, screen_y, indexing='ij')

    # TODO In @usualgaussnumber's code,
    # the flux_bndy was created when he selected a
    # region from the Matlab environment. Here, instead, we will
    # create the coordinates from the mask.
    # I'm not quite sure if this is used beyond simply blurring the
    # edges of the flux, as in pradreconstructionfluxdistsetup.m --@jtlaune

    # Scale the plasma coordinates according to the distances.
    plasma_x = screen_x/magnify
    plasma_y = screen_y/magnify

    return(plasma_x, plasma_y, flux0, flux)

def solve(X, Y, flux0, flux, dt, tol,
          chk=False, interval=1000,
          nan_exception=True, start_chk=False,
          start_phin=None, start_step=None,
          save_dir=None,
          Nstep=50000):
    """
    Main solving algorithm of the
    Monge-Ampere deflection-field potential problem using
    a finite difference scheme.

    Args:
        X (array): Plasma x-coordinates (cm).
        Y (array): Plasma y-coordinates (cm).
        flux (array): Proton flux image.
        flux0 (array): Initial flux image.
        dt (float): Timestep for Euler method.
        tol: Relaxation constant for Euler method.
        chk=False (bool): Save checkpoint files to output text files at a
                          timestep interval specified by interval.
        interval=1000 (int): Interval with which to save output files.
        nan_exception=True (bool): Optional. Set to true if you would like the
                                   algorithm to ignore NaN values.
        start_chk=False (bool): Start the solver from a specified
                                     checkpoint file.
        start_phin=None (str): 'phin' checkpoint file to use. Must be specified
                               if save_chk=True.
        start_step=None (int): Starting step of the 'phin' file. Must be
                               specified if save_chk=True.
        save_dir=None (str): If chk=True, this specifies the directory to save
                             the checkpoint files to. If this is left as None,
                             the default is the current directory.
        Nstep=50000 (int): Maximum number of steps before stopping.

    Returns:
        phin (numpy.ndarray): Screen mapping potential in the plasma
                              coordinates.
        phix (numpy.ndarray): x-gradient of the screen mapping potential in
                              the plasma coordinates.
        phiy (numpy.ndarray): x-gradient of the screen mapping potential in
                              the plasma coordinates.
    """
    plasma_x = np.copy(X)
    plasma_y = np.copy(Y)
    flux0 = np.copy(flux0)
    flux = np.copy(flux)
    flux_mean = np.mean(flux)

    if start_chk:
        start_phin = os.path.abspath(start_phin)
    if save_dir is not None:
        save_dir = os.path.abspath(save_dir)

    # Make sure that the files/directories specified are valid.
    if start_chk and (start_phin == None or start_step == None):
        raise Exception('Please specify a checkpoint phin file and '
                        'starting step.')
    if start_chk:
        if not os.path.isfile(start_phin):
            raise Exception('Please specify a valid starting checkpoint file '
                            'for phin.')
        if type(start_step) != int:
            raise Exception('Please provide a valid starting step.')

    if save_dir is not None:
        if not os.path.isdir(save_dir):
            raise Exception('Please specify a valid save directory')

    ############################
    # BEGINNING OF MAIN SOLVER #
    ############################
    if start_chk:
        # If we are starting from a checkpoint file then we load it into phin to
        # iterate.
        phin = np.loadtxt(start_phin, delimiter=',')
    if not start_chk:
        # Set up plasma parameters.
        # Initial coordinate perturbation.
        phin = .5 * (plasma_x**2 + plasma_y**2)

    # Read the size of the grid
    plasma_Nx, plasma_Ny = plasma_x.shape

    # Create plasma parameters
    dx = plasma_x[1,0] - plasma_x[0,0]
    dy = plasma_y[0,1] - plasma_y[0,0]

    plasma_Lx = plasma_x[-1,0] - plasma_x[0,0]
    plasma_Ly = plasma_y[0,-1] - plasma_y[0,0]

    # Create the derivative arrays phix, phiy, phixx, phiyy, phixy.
    # They will be the same size as phin and will be replaced every
    # timestep.
    phix = np.zeros((plasma_Nx, plasma_Ny))
    phiy = np.zeros((plasma_Nx, plasma_Ny))
    phixx = np.zeros((plasma_Nx, plasma_Ny))
    phiyy = np.zeros((plasma_Nx, plasma_Ny))
    phixy = np.zeros((plasma_Nx, plasma_Ny))

    RMSfluxerror = np.zeros((Nstep))
    pointfluxerror = np.zeros((Nstep))
    Nxrand = int(np.floor(np.random.ranf()*(plasma_Nx+1)))
    Nyrand = int(np.floor(np.random.ranf()*(plasma_Ny+1)))

    if chk:
        if save_dir==None:
            # Save the files to the current directory.
            np.savetxt(os.path.join(os.getcwd(), 'plasma_x.txt'),
                       plasma_x, delimiter=',')
            np.savetxt(os.path.join(os.getcwd(), 'plasma_y.txt'),
                       plasma_y, delimiter=',')
        else:
            # Save the files to the specified save directory.
            np.savetxt(os.path.join(save_dir, 'plasma_x.txt'),
                       plasma_x, delimiter=',')
            np.savetxt(os.path.join(save_dir, 'plasma_y.txt'),
                       plasma_y, delimiter=',')

    # Create the interpolator function.
    flux_interp = sp.interpolate.RegularGridInterpolator(
                                    (plasma_x[:,0], plasma_y[0,:]),
                                    flux,
                                    method='linear',
                                    bounds_error=False,
                                    fill_value=flux_mean)

    ###################################
    # Begin finite difference scheme. #
    ###################################
    #
    # loop Nstep times
    # 1. Check for instabilities in phin
    # 2. Create arrays for phix, phiy, phixx, phiyy, phixy
    #    - These are the derivative functions
    # 3. Finite difference scheme in 9 regions:
    #    - corners (4), edges(4), middle
    #    - get the first and second derivatives in each of these regions
    # 4. Calculate quantities and update phin
    #

    for timestep in range(Nstep):
        if start_chk:
            # If we are starting from a checkpoint file then we want to make
            # sure the timestep is accurate. Note that this will overwrite any
            # other files made after the checkpoint file we are using.
            timestep = start_step + timestep + 1

        print(timestep)

        # Test whether algorithm is stable
        if nan_exception and np.any(np.isnan(phin)):
            raise Exception('The reconstruction algorithm is unstable.')

        # Create gradient arrays phix, phiy
        phix, phiy = np.gradient(phin, dx, dy)

        # Begin adjusting gradient arrays for Neumann BCs.

        # phix corners.
        #phix[0,0] = plasma_x[0,0]
        #phix[0,plasma_Ny-1] = plasma_x[0,plasma_Ny-1]
        #phix[plasma_Nx-1,0] = plasma_x[plasma_Nx-1,0]
        #phix[plasma_Nx-1,plasma_Ny-1] = plasma_x[plasma_Nx-1,plasma_Ny-1]
        #
        ## phiy corners.
        #phiy[0,0] = plasma_y[0,0]
        #phiy[0,plasma_Ny-1] = plasma_y[0,plasma_Ny-1]
        #phiy[plasma_Nx-1,0]  = plasma_y[plasma_Nx-1,0]
        #phiy[plasma_Nx-1,plasma_Ny-1] = plasma_y[plasma_Nx-1,plasma_Ny-1]

        # phix edges.
        phix[0,:] = -plasma_x[0,:]
        phix[plasma_Nx-1,:] = plasma_x[plasma_Nx-1,:]

        # phiy edges.
        phiy[:,0] = -plasma_y[:,0]
        phiy[:,plasma_Ny-1] = plasma_y[:,plasma_Ny-1]

        # Lets try to just match up the coordinates. Seems to be a problem on
        # the boundary with the non neumann BC sides. --@jtlaune
        #phix[0,:] = plasma_x[0,:]
        #phix[-1,:] = plasma_x[-1,:]
        #phix[:,0] = plasma_x[:,0]
        #phix[:,-1] = plasma_x[:,-1]
        #
        #phiy[0,:] = plasma_y[0,:]
        #phiy[-1,:] = plasma_y[-1,:]
        #phiy[:,0] = plasma_y[:,0]
        #phiy[:,-1] = plasma_y[:,-1]

        # Create gradient arrays phixx, phiyy, phixy after we have set
        # BCs for phix and phiy.
        phixx, phixy = np.gradient(phix, dx, dy)
        _, phiyy = np.gradient(phiy, dx, dy)

        # Boundary conditions for phixy.
        # TODO Not sure what this is?? -- @jtlaune
        phixy[0,:] = 0
        phixy[plasma_Nx-1,:] = 0
        phixy[:,0] = 0
        phixy[:,plasma_Ny-1] = 0

        # Now calculate the derived quantities.
        det = phixx*phiyy - phixy**2

        # Interpolate the flux at initial points.
        # We fill points outside of the grid with flux_mean since that is what
        # we are assuming the protons are doing outside of the interaction
        # region.
        fluxn = flux_interp((phix,phiy))

        # Calculate derivative quantity.
        # TODO @usualgaussnumber's code had absolute det? Not sure why. This is
        # not mentioned in the Sulman numerical solution paper.
        Fn = np.log(fluxn*np.absolute(det)/flux0)

        # Create new phi_derived array to do error checking.
        phi_derived = phin+dt*Fn

        # Calculate the flux error.
        pointfluxerror[timestep] = (phi_derived[Nxrand,Nyrand]
                                    - phin[Nxrand,Nyrand])

        RMSfluxerror[timestep] = np.sqrt(np.mean((phi_derived-phin)**2))

        # Finally update phin for the next timestep.
        phin = phi_derived

        # TODO figure out a better way to do the tolerance checking for steady
        # state solution.
        if np.linalg.norm(Fn) < tol:
            break

        if chk:
            if (timestep % interval) == 0:
                if save_dir==None:
                    # Save the files to the current directory.
                    np.savetxt('phin' + str(timestep) + '.txt',
                               phin, delimiter=',')
                    np.savetxt('phix' + str(timestep) + '.txt',
                               phix, delimiter=',')
                    np.savetxt('phiy' + str(timestep) + '.txt',
                               phiy, delimiter=',')
                else:
                    # Save the files to the specified save directory.
                    np.savetxt(save_dir+'/phin' + str(timestep) + '.txt',
                               phin, delimiter=',')
                    np.savetxt(save_dir+'/phix' + str(timestep) + '.txt',
                               phix, delimiter=',')
                    np.savetxt(save_dir+'/phiy' + str(timestep) + '.txt',
                               phiy, delimiter=',')

        #################################
        # End finite difference scheme. #
        #################################

    return(phin, phix, phiy)
<commit_msg>Fix boundary conditions in screen.py<commit_after>"""
Main solving algorithm of the
Monge-Ampere screen mapping potential problem using
a finite difference scheme.
"""
import numpy as np
import scipy as sp
import scipy.interpolate
import os.path

def setup(ri, li, rs, dxS, image, mask):
    """
    Takes proton radiography distance parameters, a proton flux image, and a
    flux image selection mask and sets up the deflection-field
    potential problem.

    Args:
        ri (float): Distance from source to plasma (cm).
        li (float): Distance across plasma (cm).
        rs (float): Distance from plasma to screen (cm).
        dxS (float): Pixel size (cm)
        image (array): Proton flux image to analyze.
        mask (array): Bitwise mask of data to analyze in the flux image,
                      assuming 1 is masked and 0 is unmasked.

    Returns:
        plasma_x (array): Plasma x-coordinates (cm).
        plasma_y (array): Plasma x-coordinates (cm).
        flux0 (array): Initial flux image.
        flux (array): Flux image.
    """

    magnify = (rs+ri+.5*li)/(ri+.5*li)
    # Leave the original image unchanged.
    flux_image = np.copy(image)
    flux_mask = np.copy(mask)

    # Mask the flux_image and
    # fill the masked positions with the mean flux.
    #
    # TODO This may add gradient to the proton fluence and therefore
    # create artificial magnetic fields-- may need to find a different way to
    # do this. Possibly extrapolate gradients from the boundary?
    #flux = np.ma.array(flux_image, mask=np.logical_not(flux_mask))
    flux = flux_image * np.logical_not(flux_mask)
    flux_mean = np.nanmean(flux)
    flux = flux + flux_mean * flux_mask

    # Calculate the total flux throughout the selected region
    flux_tot = np.nansum(np.logical_not(flux_mask)*flux)

    # TODO
    # @usualgaussnumber applies a Gaussian filter after this step, see
    # pradreconstructionfluxdistsetup.m

    # Prepare the flux and
    # flux_0 for the primary reconstruction process.
    # flux_0 and flux_0 mask are the initial flux distributions,
    # which we choose to be uniform over the selected region while
    # preserving total flux.
    # Then add a shadow flux region to flux0 which is equal to the flux mean.
    flux0_mask = np.copy(flux_mask)
    flux0 = np.zeros(flux.shape)
    flux0 = (flux_tot * np.logical_not(flux0_mask))/(np.nansum(np.logical_not(flux0_mask)))
    flux0 = flux0 + flux0_mask * np.mean(flux0)

    # Variables:
    # Nx, N1 are the number of cells on the 0th and 1st axis
    # dx is the length of each cell (square)
    # L0, L1 are the lengths of the flux image in the 0th and 1st axis
    Nx, Ny = flux.shape
    Lx = dxS*Nx
    Ly = dxS*Ny

    # Create the coordinates for the flux samples for the centers of the
    # bins on the screen.
    screen_x = np.linspace((-Lx/2)+.5*dxS, (Lx/2)-.5*dxS, Nx)
    screen_y = np.linspace((-Ly/2)+.5*dxS, (Ly/2)-.5*dxS, Ny)
    screen_x, screen_y = np.meshgrid(screen_x, screen_y, indexing='ij')

    # TODO In @usualgaussnumber's code,
    # the flux_bndy was created when he selected a
    # region from the Matlab environment. Here, instead, we will
    # create the coordinates from the mask.
    # I'm not quite sure if this is used beyond simply blurring the
    # edges of the flux, as in pradreconstructionfluxdistsetup.m --@jtlaune

    # Scale the plasma coordinates according to the distances.
    plasma_x = screen_x/magnify
    plasma_y = screen_y/magnify

    return(plasma_x, plasma_y, flux0, flux)

def solve(X, Y, flux0, flux, dt, tol,
          chk=False, interval=1000,
          nan_exception=True, start_chk=False,
          start_phin=None, start_step=None,
          save_dir=None,
          Nstep=50000):
    """
    Main solving algorithm of the
    Monge-Ampere deflection-field potential problem using
    a finite difference scheme.

    Args:
        X (array): Plasma x-coordinates (cm).
        Y (array): Plasma y-coordinates (cm).
        flux (array): Proton flux image.
        flux0 (array): Initial flux image.
        dt (float): Timestep for Euler method.
        tol: Relaxation constant for Euler method.
        chk=False (bool): Save checkpoint files to output text files at a
                          timestep interval specified by interval.
        interval=1000 (int): Interval with which to save output files.
        nan_exception=True (bool): Optional. Set to true if you would like the
                                   algorithm to ignore NaN values.
        start_chk=False (bool): Start the solver from a specified
                                     checkpoint file.
        start_phin=None (str): 'phin' checkpoint file to use. Must be specified
                               if save_chk=True.
        start_step=None (int): Starting step of the 'phin' file. Must be
                               specified if save_chk=True.
        save_dir=None (str): If chk=True, this specifies the directory to save
                             the checkpoint files to. If this is left as None,
                             the default is the current directory.
        Nstep=50000 (int): Maximum number of steps before stopping.

    Returns:
        phin (numpy.ndarray): Screen mapping potential in the plasma
                              coordinates.
        phix (numpy.ndarray): x-gradient of the screen mapping potential in
                              the plasma coordinates.
        phiy (numpy.ndarray): x-gradient of the screen mapping potential in
                              the plasma coordinates.
    """
    plasma_x = np.copy(X)
    plasma_y = np.copy(Y)
    flux0 = np.copy(flux0)
    flux = np.copy(flux)
    flux_mean = np.mean(flux)

    if start_chk:
        start_phin = os.path.abspath(start_phin)
    if save_dir is not None:
        save_dir = os.path.abspath(save_dir)

    # Make sure that the files/directories specified are valid.
    if start_chk and (start_phin == None or start_step == None):
        raise Exception('Please specify a checkpoint phin file and '
                        'starting step.')
    if start_chk:
        if not os.path.isfile(start_phin):
            raise Exception('Please specify a valid starting checkpoint file '
                            'for phin.')
        if type(start_step) != int:
            raise Exception('Please provide a valid starting step.')

    if save_dir is not None:
        if not os.path.isdir(save_dir):
            raise Exception('Please specify a valid save directory')

    ############################
    # BEGINNING OF MAIN SOLVER #
    ############################
    if start_chk:
        # If we are starting from a checkpoint file then we load it into phin to
        # iterate.
        phin = np.loadtxt(start_phin, delimiter=',')
    if not start_chk:
        # Set up plasma parameters.
        # Initial coordinate perturbation.
        phin = .5 * (plasma_x**2 + plasma_y**2)

    # Read the size of the grid
    plasma_Nx, plasma_Ny = plasma_x.shape

    # Create plasma parameters
    dx = plasma_x[1,0] - plasma_x[0,0]
    dy = plasma_y[0,1] - plasma_y[0,0]

    plasma_Lx = plasma_x[-1,0] - plasma_x[0,0]
    plasma_Ly = plasma_y[0,-1] - plasma_y[0,0]

    # Create the derivative arrays phix, phiy, phixx, phiyy, phixy.
    # They will be the same size as phin and will be replaced every
    # timestep.
    phix = np.zeros((plasma_Nx, plasma_Ny))
    phiy = np.zeros((plasma_Nx, plasma_Ny))
    phixx = np.zeros((plasma_Nx, plasma_Ny))
    phiyy = np.zeros((plasma_Nx, plasma_Ny))
    phixy = np.zeros((plasma_Nx, plasma_Ny))

    RMSfluxerror = np.zeros((Nstep))
    pointfluxerror = np.zeros((Nstep))
    Nxrand = int(np.floor(np.random.ranf()*(plasma_Nx+1)))
    Nyrand = int(np.floor(np.random.ranf()*(plasma_Ny+1)))

    if chk:
        if save_dir==None:
            # Save the files to the current directory.
            np.savetxt(os.path.join(os.getcwd(), 'plasma_x.txt'),
                       plasma_x, delimiter=',')
            np.savetxt(os.path.join(os.getcwd(), 'plasma_y.txt'),
                       plasma_y, delimiter=',')
        else:
            # Save the files to the specified save directory.
            np.savetxt(os.path.join(save_dir, 'plasma_x.txt'),
                       plasma_x, delimiter=',')
            np.savetxt(os.path.join(save_dir, 'plasma_y.txt'),
                       plasma_y, delimiter=',')

    # Create the interpolator function.
    flux_interp = sp.interpolate.RegularGridInterpolator(
                                    (plasma_x[:,0], plasma_y[0,:]),
                                    flux,
                                    method='linear',
                                    bounds_error=False,
                                    fill_value=flux_mean)

    ###################################
    # Begin finite difference scheme. #
    ###################################
    #
    # loop Nstep times
    # 1. Check for instabilities in phin
    # 2. Create arrays for phix, phiy, phixx, phiyy, phixy
    #    - These are the derivative functions
    # 3. Finite difference scheme in 9 regions:
    #    - corners (4), edges(4), middle
    #    - get the first and second derivatives in each of these regions
    # 4. Calculate quantities and update phin
    #

    for timestep in range(Nstep):
        if start_chk:
            # If we are starting from a checkpoint file then we want to make
            # sure the timestep is accurate. Note that this will overwrite any
            # other files made after the checkpoint file we are using.
            timestep = start_step + timestep + 1

        print(timestep)

        # Test whether algorithm is stable
        if nan_exception and np.any(np.isnan(phin)):
            raise Exception('The reconstruction algorithm is unstable.')

        # Create gradient arrays phix, phiy
        phix, phiy = np.gradient(phin, dx, dy)

        # Begin adjusting gradient arrays for Neumann BCs.

        # phix corners.
        #phix[0,0] = plasma_x[0,0]
        #phix[0,plasma_Ny-1] = plasma_x[0,plasma_Ny-1]
        #phix[plasma_Nx-1,0] = plasma_x[plasma_Nx-1,0]
        #phix[plasma_Nx-1,plasma_Ny-1] = plasma_x[plasma_Nx-1,plasma_Ny-1]
        #
        ## phiy corners.
        #phiy[0,0] = plasma_y[0,0]
        #phiy[0,plasma_Ny-1] = plasma_y[0,plasma_Ny-1]
        #phiy[plasma_Nx-1,0]  = plasma_y[plasma_Nx-1,0]
        #phiy[plasma_Nx-1,plasma_Ny-1] = plasma_y[plasma_Nx-1,plasma_Ny-1]

        # phix edges.
        phix[0,:] = plasma_x[0,:]
        phix[plasma_Nx-1,:] = plasma_x[plasma_Nx-1,:]

        # phiy edges.
        phiy[:,0] = plasma_y[:,0]
        phiy[:,plasma_Ny-1] = plasma_y[:,plasma_Ny-1]

        # Lets try to just match up the coordinates. Seems to be a problem on
        # the boundary with the non neumann BC sides. --@jtlaune
        #phix[0,:] = plasma_x[0,:]
        #phix[-1,:] = plasma_x[-1,:]
        #phix[:,0] = plasma_x[:,0]
        #phix[:,-1] = plasma_x[:,-1]
        #
        #phiy[0,:] = plasma_y[0,:]
        #phiy[-1,:] = plasma_y[-1,:]
        #phiy[:,0] = plasma_y[:,0]
        #phiy[:,-1] = plasma_y[:,-1]

        # Create gradient arrays phixx, phiyy, phixy after we have set
        # BCs for phix and phiy.
        phixx, phixy = np.gradient(phix, dx, dy)
        _, phiyy = np.gradient(phiy, dx, dy)

        # Boundary conditions for phixy.
        # TODO Not sure what this is?? -- @jtlaune
        phixy[0,:] = 0
        phixy[plasma_Nx-1,:] = 0
        phixy[:,0] = 0
        phixy[:,plasma_Ny-1] = 0

        # Now calculate the derived quantities.
        det = phixx*phiyy - phixy**2

        # Interpolate the flux at initial points.
        # We fill points outside of the grid with flux_mean since that is what
        # we are assuming the protons are doing outside of the interaction
        # region.
        fluxn = flux_interp((phix,phiy))

        # Calculate derivative quantity.
        # TODO @usualgaussnumber's code had absolute det? Not sure why. This is
        # not mentioned in the Sulman numerical solution paper.
        Fn = np.log(fluxn*np.absolute(det)/flux0)

        # Create new phi_derived array to do error checking.
        phi_derived = phin+dt*Fn

        # Calculate the flux error.
        pointfluxerror[timestep] = (phi_derived[Nxrand,Nyrand]
                                    - phin[Nxrand,Nyrand])

        RMSfluxerror[timestep] = np.sqrt(np.mean((phi_derived-phin)**2))

        # Finally update phin for the next timestep.
        phin = phi_derived

        # TODO figure out a better way to do the tolerance checking for steady
        # state solution.
        if np.linalg.norm(Fn) < tol:
            break

        if chk:
            if (timestep % interval) == 0:
                if save_dir==None:
                    # Save the files to the current directory.
                    np.savetxt('phin' + str(timestep) + '.txt',
                               phin, delimiter=',')
                    np.savetxt('phix' + str(timestep) + '.txt',
                               phix, delimiter=',')
                    np.savetxt('phiy' + str(timestep) + '.txt',
                               phiy, delimiter=',')
                else:
                    # Save the files to the specified save directory.
                    np.savetxt(save_dir+'/phin' + str(timestep) + '.txt',
                               phin, delimiter=',')
                    np.savetxt(save_dir+'/phix' + str(timestep) + '.txt',
                               phix, delimiter=',')
                    np.savetxt(save_dir+'/phiy' + str(timestep) + '.txt',
                               phiy, delimiter=',')

        #################################
        # End finite difference scheme. #
        #################################

    return(phin, phix, phiy)
<|endoftext|>