<commit_before>from enum import Enum
from typing import OrderedDict

import numpy as np
import taichi.core.primitive_types as primitive_type
from taichi.core.util import ti_core as _ti_core
from taichi.lang import impl
from taichi.lang.enums import Layout
from taichi.lang.exception import InvalidOperationError, TaichiSyntaxError
from taichi.lang.field import Field, ScalarField
from taichi.lang.kernel_arguments import template
from taichi.lang.kernel_impl import kernel
from taichi.lang.matrix import Matrix, MatrixField
from taichi.lang.snode import SNode
from taichi.lang.struct import Struct, StructField
from taichi.lang.types import CompoundType
from taichi.lang.util import (cook_dtype, has_pytorch, is_taichi_class,
                              python_scope, taichi_scope, to_pytorch_type)
from taichi.misc.mesh_loader import *

import taichi as ti

MeshTopology = _ti_core.MeshTopology
MeshElementType = _ti_core.MeshElementType
MeshRelationType = _ti_core.MeshRelationType
MeshElementReorderingType = _ti_core.MeshElementReorderingType
ConvType = _ti_core.ConvType
element_order = _ti_core.element_order
from_end_element_order = _ti_core.from_end_element_order
to_end_element_order = _ti_core.to_end_element_order
relation_by_orders = _ti_core.relation_by_orders
inverse_relation = _ti_core.inverse_relation

class MeshAttrType:
    def __init__(self, name, dtype, reordering, needs_grad):
        self.name = name
        self.dtype = dtype
        self.reordering = reordering
        self.needs_grad = needs_grad

class MeshElementField:
    def __init__(self, mesh_instance, type, attr_dict, field_dict):
        self.mesh = mesh_instance
        self.type = type
        self.attr_dict = attr_dict
        self.field_dict = field_dict

        self.register_fields()

    @property
    def keys(self):
        return list(self.field_dict.keys())

    @property
    def members(self):
        return list(self.field_dict.values())

    @property
    def items(self):
        return self.field_dict.items()

    @staticmethod
    def make_getter(key):
        def getter(self):
            """Get an entry from custom struct by name."""
            _taichi_skip_traceback = 1
            return self.field_dict[key]

        return getter

    @staticmethod
    def make_setter(key):
        @python_scope
        def setter(self, value):
            _taichi_skip_traceback = 1
            self.field_dict[key] = value

        return setter

    def register_fields(self):
        for v in self.get_field_members():
            _ti_core.insert_mesh_attribute(self.mesh.mesh_ptr, self.type,
                                           v.ptr)

        for k in self.keys:
            setattr(
                MeshElementField, k,
                property(
                    MeshElementField.make_getter(k),
                    MeshElementField.make_setter(k),
                ))

    def get_field_members(self):
        field_members = []
        for m in self.members:
            assert isinstance(m, Field)
            field_members += m.get_field_members()
        return field_members

    @property
    def snode(self):
        return self.members[0].snode

    def loop_range(self):
        return self.members[0].loop_range()

    @python_scope
    def copy_from(self, other):
        assert isinstance(other, Field)
        assert set(self.keys) == set(other.keys)
        for k in self.keys:
            self.field_dict[k].copy_from(other[k])

    @python_scope
    def fill(self, val):
        for v in self.members:
            v.fill(val)

    def initialize_host_accessors(self):
        for v in self.members:
            v.initialize_host_accessors()

    def get_member_field(self, key):
        return self.field_dict[key]

    @python_scope
    def from_numpy(self, array_dict):
        for k, v in self.items:
            v.from_numpy(array_dict[k])

    @python_scope
    def from_torch(self, array_dict):
        for k, v in self.items:
            v.from_torch(array_dict[k])

    @python_scope
    def to_numpy(self):
        return {k: v.to_numpy() for k, v in self.items}

    @python_scope
    def to_torch(self, device=None):
        return {k: v.to_torch(device=device) for k, v in self.items}

    @python_scope
    def __setitem__(self, indices,
                    element):  # TODO(changyu): handle reordering case
        self.initialize_host_accessors()
        self[indices].set_entries(element)

    @python_scope
    def __getitem__(self, indices):  # TODO(changyu): handle reordering case
        self.initialize_host_accessors()
        # scalar fields does not instantiate SNodeHostAccess by default
        entries = {
            k: v.host_access(self.pad_key(indices))[0] if isinstance(
                v, ScalarField) else v[indices]
            for k, v in self.items
        }
        return Struct(entries)

class MeshElement:
    def __init__(self, type):
        self.type = type
        self.layout = Layout.SOA
        self.attr_dict = {}

    def _SOA(self, soa=True):  # AOS/SOA
        self.layout = Layout.SOA if soa else Layout.AOS

    def _AOS(self, aos=True):
        self.layout = Layout.AOS if aos else Layout.SOA

    SOA = property(fget=_SOA)
    AOS = property(fget=_AOS)

    def place(
        self,
        members,
        reordering=MeshElementReorderingType.NonReordering,
        needs_grad=False,
    ):
        for key, dtype in members.items():
            self.attr_dict[key] = MeshAttrType(key, dtype, reordering,
                                               needs_grad)

    def build(self, mesh_instance, size):
        field_dict = {}

        for key, attr in self.attr_dict.items():
            if isinstance(attr.dtype, CompoundType):
                field_dict[key] = attr.dtype.field(shape=None,
                                                   needs_grad=attr.needs_grad)
            else:
                field_dict[key] = impl.field(attr.dtype,
                                             shape=None,
                                             needs_grad=attr.needs_grad)

        if self.layout == Layout.SOA:
            for key, field in field_dict.items():
                impl.root.dense(impl.indices(0), size).place(field)
                if self.attr_dict[key].needs_grad:
                    impl.root.dense(impl.indices(0), size).place(field.grad)
        else:
            impl.root.dense(impl.indices(0),
                            size).place(*tuple(field_dict.values()))
            grads = []
            for key, field in field_dict.items():
                if self.attr_dict[key].needs_gard: grads.append(field.grad)
            impl.root.dense(impl.indices(0), size).place(*grads)

        return MeshElementField(mesh_instance, self.type, self.attr_dict,
                                field_dict)

# Define the instance of the Mesh Type, stores the field (type and data) info
class MeshInstance:
    def __init__(self, type):
        self.type = type
        self.mesh_ptr = _ti_core.create_mesh()

    def set_owned_offset(self, element_type: MeshElementType,
                         owned_offset: ScalarField):
        _ti_core.set_owned_offset(self.mesh_ptr, element_type,
                                  owned_offset.vars[0].ptr.snode())

    def set_total_offset(self, element_type: MeshElementType,
                         total_offset: ScalarField):
        _ti_core.set_total_offset(self.mesh_ptr, element_type,
                                  total_offset.vars[0].ptr.snode())

    def set_l2g(self, element_type: MeshElementType,
                total_offset: ScalarField):
        _ti_core.set_l2g(self.mesh_ptr, element_type,
                         total_offset.vars[0].ptr.snode())

    def set_l2r(self, element_type: MeshElementType,
                total_offset: ScalarField):
        _ti_core.set_l2r(self.mesh_ptr, element_type,
                         total_offset.vars[0].ptr.snode())

    def set_num_patches(self, num_patches: int):
        _ti_core.set_num_patches(self.mesh_ptr, num_patches)

    def set_patch_max_element_num(self, element_type: MeshElementType,
                                  max_element_num: int):
        _ti_core.set_patch_max_element_num(self.mesh_ptr, element_type,
                                           max_element_num)

    def set_relation_fixed(self, rel_type: MeshRelationType,
                           value: ScalarField):
        _ti_core.set_relation_fixed(self.mesh_ptr, rel_type,
                                    value.vars[0].ptr.snode())

    def set_relation_dynamic(self, rel_type: MeshRelationType,
                             value: ScalarField, offset: ScalarField):
        _ti_core.set_relation_dynamic(self.mesh_ptr, rel_type,
                                      value.vars[0].ptr.snode(),
                                      offset.vars[0].ptr.snode())

    def add_mesh_attribute(self, element_type, snode, reordering_type):
        _ti_core.add_mesh_attribute(self.mesh_ptr, element_type, snode,
                                    reordering_type)

# Define the Mesh Type, stores the field type info
class MeshType:
    def __init__(self, topology):
        if not ti.is_extension_supported(ti.cfg.arch, ti.extension.mesh):
            raise Exception('Backend ' + str(ti.cfg.arch) +
                            ' doesn\'t support MeshTaichi extension')

        self.topology = topology
        self.verts = MeshElement(MeshElementType.Vertex)
        self.edges = MeshElement(MeshElementType.Edge)
        self.faces = MeshElement(MeshElementType.Face)
        if topology == MeshTopology.Tetrahedron:
            self.cells = MeshElement(MeshElementType.Cell)

    def build(self, size: tuple):
        instance = MeshInstance(self)
        instance.fields = {}
        if size[0] > 0:
            instance.verts = self.verts.build(instance, size[0])
            instance.fields[MeshElementType.Vertex] = instance.verts
        if size[1] > 0:
            instance.edges = self.edges.build(instance, size[1])
            instance.fields[MeshElementType.Edge] = instance.edges
        if size[2] > 0:
            instance.faces = self.faces.build(instance, size[2])
            instance.fields[MeshElementType.Face] = instance.faces
        if self.topology == MeshTopology.Tetrahedron and size[3] > 0:
            instance.cells = self.cells.build(instance, size[3])
            instance.fields[MeshElementType.Cell] = instance.cells
        return instance

# Mesh First Class
class Mesh:
    def __init__(self):
        pass

    non_reordering = MeshElementReorderingType.NonReordering
    reordering = MeshElementReorderingType.Reordering
    surface_first = MeshElementReorderingType.SurfaceFirst
    cell_first = MeshElementReorderingType.CellFirst

    @staticmethod
    def Tet():
        return MeshType(MeshTopology.Tetrahedron)

    @staticmethod
    def Tri():
        return MeshType(MeshTopology.Triangle)

def TriMesh():
    return Mesh.Tri()

def TetMesh():
    return Mesh.Tet()

class MeshElementFieldProxy:
    def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
                 entry_expr: impl.Expr):
        self.mesh = mesh
        self.element_type = element_type
        self.entry_expr = entry_expr

        element_field = self.mesh.fields[self.element_type]
        for key, attr in element_field.field_dict.items():
            global_entry_expr = impl.Expr(
                _ti_core.get_index_conversion(
                    self.mesh.mesh_ptr, entry_expr,
                    ConvType.l2g if element_field.attr_dict[key].reordering
                    == MeshElementReorderingType.NonReordering else
                    ConvType.l2r))  # transform index space
            global_entry_expr_group = impl.make_expr_group(
                *tuple([global_entry_expr]))
            if isinstance(attr, MatrixField):
                setattr(
                    self, key,
                    Matrix.with_entries(attr.n, attr.m, [
                        impl.Expr(
                            _ti_core.subscript(e.ptr, global_entry_expr_group))
                        for e in attr.get_field_members()
                    ]))
            elif isinstance(attr, StructField):
                raise RuntimeError('ti.Mesh has not support StructField yet')
            else:  # isinstance(attr, Field)
                var = attr.get_field_members()[0].ptr
                setattr(
                    self, key,
                    impl.Expr(_ti_core.subscript(var,
                                                 global_entry_expr_group)))

    @property
    def ptr(self):
        return self.entry_expr

    @property
    def id(self):  # return the global non-reordered index
        l2g_expr = impl.Expr(
            _ti_core.get_index_conversion(self.mesh.mesh_ptr, self.entry_expr,
                                          ConvType.l2g))
        return l2g_expr

class MeshRelationAccessProxy:
    def __init__(self, mesh: MeshInstance, from_index: impl.Expr,
                 to_element_type: MeshElementType):
        self.mesh = mesh
        self.from_index = from_index
        self.to_element_type = to_element_type

    @property
    def size(self):
        return impl.Expr(
            _ti_core.get_relation_size(self.mesh.mesh_ptr, self.from_index.ptr,
                                       self.to_element_type))

    def subscript(self, *indices):
        assert (len(indices) == 1)
        entry_expr = _ti_core.get_relation_access(self.mesh.mesh_ptr,
                                                  self.from_index.ptr,
                                                  self.to_element_type,
                                                  impl.Expr(indices[0]).ptr)
        return MeshElementFieldProxy(self.mesh, self.to_element_type,
                                     entry_expr)
<commit_msg>[Mesh] Use ti.axes instead of ti.indices<commit_after>from enum import Enum
from typing import OrderedDict

import numpy as np
import taichi.core.primitive_types as primitive_type
from taichi.core.util import ti_core as _ti_core
from taichi.lang import impl
from taichi.lang.enums import Layout
from taichi.lang.exception import InvalidOperationError, TaichiSyntaxError
from taichi.lang.field import Field, ScalarField
from taichi.lang.kernel_arguments import template
from taichi.lang.kernel_impl import kernel
from taichi.lang.matrix import Matrix, MatrixField
from taichi.lang.snode import SNode
from taichi.lang.struct import Struct, StructField
from taichi.lang.types import CompoundType
from taichi.lang.util import (cook_dtype, has_pytorch, is_taichi_class,
                              python_scope, taichi_scope, to_pytorch_type)
from taichi.misc.mesh_loader import *

import taichi as ti

MeshTopology = _ti_core.MeshTopology
MeshElementType = _ti_core.MeshElementType
MeshRelationType = _ti_core.MeshRelationType
MeshElementReorderingType = _ti_core.MeshElementReorderingType
ConvType = _ti_core.ConvType
element_order = _ti_core.element_order
from_end_element_order = _ti_core.from_end_element_order
to_end_element_order = _ti_core.to_end_element_order
relation_by_orders = _ti_core.relation_by_orders
inverse_relation = _ti_core.inverse_relation

class MeshAttrType:
    def __init__(self, name, dtype, reordering, needs_grad):
        self.name = name
        self.dtype = dtype
        self.reordering = reordering
        self.needs_grad = needs_grad

class MeshElementField:
    def __init__(self, mesh_instance, type, attr_dict, field_dict):
        self.mesh = mesh_instance
        self.type = type
        self.attr_dict = attr_dict
        self.field_dict = field_dict

        self.register_fields()

    @property
    def keys(self):
        return list(self.field_dict.keys())

    @property
    def members(self):
        return list(self.field_dict.values())

    @property
    def items(self):
        return self.field_dict.items()

    @staticmethod
    def make_getter(key):
        def getter(self):
            """Get an entry from custom struct by name."""
            _taichi_skip_traceback = 1
            return self.field_dict[key]

        return getter

    @staticmethod
    def make_setter(key):
        @python_scope
        def setter(self, value):
            _taichi_skip_traceback = 1
            self.field_dict[key] = value

        return setter

    def register_fields(self):
        for v in self.get_field_members():
            _ti_core.insert_mesh_attribute(self.mesh.mesh_ptr, self.type,
                                           v.ptr)

        for k in self.keys:
            setattr(
                MeshElementField, k,
                property(
                    MeshElementField.make_getter(k),
                    MeshElementField.make_setter(k),
                ))

    def get_field_members(self):
        field_members = []
        for m in self.members:
            assert isinstance(m, Field)
            field_members += m.get_field_members()
        return field_members

    @property
    def snode(self):
        return self.members[0].snode

    def loop_range(self):
        return self.members[0].loop_range()

    @python_scope
    def copy_from(self, other):
        assert isinstance(other, Field)
        assert set(self.keys) == set(other.keys)
        for k in self.keys:
            self.field_dict[k].copy_from(other[k])

    @python_scope
    def fill(self, val):
        for v in self.members:
            v.fill(val)

    def initialize_host_accessors(self):
        for v in self.members:
            v.initialize_host_accessors()

    def get_member_field(self, key):
        return self.field_dict[key]

    @python_scope
    def from_numpy(self, array_dict):
        for k, v in self.items:
            v.from_numpy(array_dict[k])

    @python_scope
    def from_torch(self, array_dict):
        for k, v in self.items:
            v.from_torch(array_dict[k])

    @python_scope
    def to_numpy(self):
        return {k: v.to_numpy() for k, v in self.items}

    @python_scope
    def to_torch(self, device=None):
        return {k: v.to_torch(device=device) for k, v in self.items}

    @python_scope
    def __setitem__(self, indices,
                    element):  # TODO(changyu): handle reordering case
        self.initialize_host_accessors()
        self[indices].set_entries(element)

    @python_scope
    def __getitem__(self, indices):  # TODO(changyu): handle reordering case
        self.initialize_host_accessors()
        # scalar fields does not instantiate SNodeHostAccess by default
        entries = {
            k: v.host_access(self.pad_key(indices))[0] if isinstance(
                v, ScalarField) else v[indices]
            for k, v in self.items
        }
        return Struct(entries)

class MeshElement:
    def __init__(self, type):
        self.type = type
        self.layout = Layout.SOA
        self.attr_dict = {}

    def _SOA(self, soa=True):  # AOS/SOA
        self.layout = Layout.SOA if soa else Layout.AOS

    def _AOS(self, aos=True):
        self.layout = Layout.AOS if aos else Layout.SOA

    SOA = property(fget=_SOA)
    AOS = property(fget=_AOS)

    def place(
        self,
        members,
        reordering=MeshElementReorderingType.NonReordering,
        needs_grad=False,
    ):
        for key, dtype in members.items():
            self.attr_dict[key] = MeshAttrType(key, dtype, reordering,
                                               needs_grad)

    def build(self, mesh_instance, size):
        field_dict = {}

        for key, attr in self.attr_dict.items():
            if isinstance(attr.dtype, CompoundType):
                field_dict[key] = attr.dtype.field(shape=None,
                                                   needs_grad=attr.needs_grad)
            else:
                field_dict[key] = impl.field(attr.dtype,
                                             shape=None,
                                             needs_grad=attr.needs_grad)

        if self.layout == Layout.SOA:
            for key, field in field_dict.items():
                impl.root.dense(impl.axes(0), size).place(field)
                if self.attr_dict[key].needs_grad:
                    impl.root.dense(impl.axes(0), size).place(field.grad)
        else:
            impl.root.dense(impl.axes(0),
                            size).place(*tuple(field_dict.values()))
            grads = []
            for key, field in field_dict.items():
                if self.attr_dict[key].needs_gard: grads.append(field.grad)
            impl.root.dense(impl.axes(0), size).place(*grads)

        return MeshElementField(mesh_instance, self.type, self.attr_dict,
                                field_dict)

# Define the instance of the Mesh Type, stores the field (type and data) info
class MeshInstance:
    def __init__(self, type):
        self.type = type
        self.mesh_ptr = _ti_core.create_mesh()

    def set_owned_offset(self, element_type: MeshElementType,
                         owned_offset: ScalarField):
        _ti_core.set_owned_offset(self.mesh_ptr, element_type,
                                  owned_offset.vars[0].ptr.snode())

    def set_total_offset(self, element_type: MeshElementType,
                         total_offset: ScalarField):
        _ti_core.set_total_offset(self.mesh_ptr, element_type,
                                  total_offset.vars[0].ptr.snode())

    def set_l2g(self, element_type: MeshElementType,
                total_offset: ScalarField):
        _ti_core.set_l2g(self.mesh_ptr, element_type,
                         total_offset.vars[0].ptr.snode())

    def set_l2r(self, element_type: MeshElementType,
                total_offset: ScalarField):
        _ti_core.set_l2r(self.mesh_ptr, element_type,
                         total_offset.vars[0].ptr.snode())

    def set_num_patches(self, num_patches: int):
        _ti_core.set_num_patches(self.mesh_ptr, num_patches)

    def set_patch_max_element_num(self, element_type: MeshElementType,
                                  max_element_num: int):
        _ti_core.set_patch_max_element_num(self.mesh_ptr, element_type,
                                           max_element_num)

    def set_relation_fixed(self, rel_type: MeshRelationType,
                           value: ScalarField):
        _ti_core.set_relation_fixed(self.mesh_ptr, rel_type,
                                    value.vars[0].ptr.snode())

    def set_relation_dynamic(self, rel_type: MeshRelationType,
                             value: ScalarField, offset: ScalarField):
        _ti_core.set_relation_dynamic(self.mesh_ptr, rel_type,
                                      value.vars[0].ptr.snode(),
                                      offset.vars[0].ptr.snode())

    def add_mesh_attribute(self, element_type, snode, reordering_type):
        _ti_core.add_mesh_attribute(self.mesh_ptr, element_type, snode,
                                    reordering_type)

# Define the Mesh Type, stores the field type info
class MeshType:
    def __init__(self, topology):
        if not ti.is_extension_supported(ti.cfg.arch, ti.extension.mesh):
            raise Exception('Backend ' + str(ti.cfg.arch) +
                            ' doesn\'t support MeshTaichi extension')

        self.topology = topology
        self.verts = MeshElement(MeshElementType.Vertex)
        self.edges = MeshElement(MeshElementType.Edge)
        self.faces = MeshElement(MeshElementType.Face)
        if topology == MeshTopology.Tetrahedron:
            self.cells = MeshElement(MeshElementType.Cell)

    def build(self, size: tuple):
        instance = MeshInstance(self)
        instance.fields = {}
        if size[0] > 0:
            instance.verts = self.verts.build(instance, size[0])
            instance.fields[MeshElementType.Vertex] = instance.verts
        if size[1] > 0:
            instance.edges = self.edges.build(instance, size[1])
            instance.fields[MeshElementType.Edge] = instance.edges
        if size[2] > 0:
            instance.faces = self.faces.build(instance, size[2])
            instance.fields[MeshElementType.Face] = instance.faces
        if self.topology == MeshTopology.Tetrahedron and size[3] > 0:
            instance.cells = self.cells.build(instance, size[3])
            instance.fields[MeshElementType.Cell] = instance.cells
        return instance

# Mesh First Class
class Mesh:
    def __init__(self):
        pass

    non_reordering = MeshElementReorderingType.NonReordering
    reordering = MeshElementReorderingType.Reordering
    surface_first = MeshElementReorderingType.SurfaceFirst
    cell_first = MeshElementReorderingType.CellFirst

    @staticmethod
    def Tet():
        return MeshType(MeshTopology.Tetrahedron)

    @staticmethod
    def Tri():
        return MeshType(MeshTopology.Triangle)

def TriMesh():
    return Mesh.Tri()

def TetMesh():
    return Mesh.Tet()

class MeshElementFieldProxy:
    def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
                 entry_expr: impl.Expr):
        self.mesh = mesh
        self.element_type = element_type
        self.entry_expr = entry_expr

        element_field = self.mesh.fields[self.element_type]
        for key, attr in element_field.field_dict.items():
            global_entry_expr = impl.Expr(
                _ti_core.get_index_conversion(
                    self.mesh.mesh_ptr, entry_expr,
                    ConvType.l2g if element_field.attr_dict[key].reordering
                    == MeshElementReorderingType.NonReordering else
                    ConvType.l2r))  # transform index space
            global_entry_expr_group = impl.make_expr_group(
                *tuple([global_entry_expr]))
            if isinstance(attr, MatrixField):
                setattr(
                    self, key,
                    Matrix.with_entries(attr.n, attr.m, [
                        impl.Expr(
                            _ti_core.subscript(e.ptr, global_entry_expr_group))
                        for e in attr.get_field_members()
                    ]))
            elif isinstance(attr, StructField):
                raise RuntimeError('ti.Mesh has not support StructField yet')
            else:  # isinstance(attr, Field)
                var = attr.get_field_members()[0].ptr
                setattr(
                    self, key,
                    impl.Expr(_ti_core.subscript(var,
                                                 global_entry_expr_group)))

    @property
    def ptr(self):
        return self.entry_expr

    @property
    def id(self):  # return the global non-reordered index
        l2g_expr = impl.Expr(
            _ti_core.get_index_conversion(self.mesh.mesh_ptr, self.entry_expr,
                                          ConvType.l2g))
        return l2g_expr

class MeshRelationAccessProxy:
    def __init__(self, mesh: MeshInstance, from_index: impl.Expr,
                 to_element_type: MeshElementType):
        self.mesh = mesh
        self.from_index = from_index
        self.to_element_type = to_element_type

    @property
    def size(self):
        return impl.Expr(
            _ti_core.get_relation_size(self.mesh.mesh_ptr, self.from_index.ptr,
                                       self.to_element_type))

    def subscript(self, *indices):
        assert (len(indices) == 1)
        entry_expr = _ti_core.get_relation_access(self.mesh.mesh_ptr,
                                                  self.from_index.ptr,
                                                  self.to_element_type,
                                                  impl.Expr(indices[0]).ptr)
        return MeshElementFieldProxy(self.mesh, self.to_element_type,
                                     entry_expr)
<|endoftext|>


<commit_before># Copyright (c) 2015, Frappe Technologies Pvt. Ltd. and Contributors
# MIT License. See license.txt

from __future__ import unicode_literals

import frappe
from frappe.utils import time_diff_in_seconds, now, now_datetime, DATETIME_FORMAT
from dateutil.relativedelta import relativedelta
from six import string_types
import json

@frappe.whitelist()
def get_notifications():
	if (frappe.flags.in_install or
		not frappe.db.get_single_value('System Settings', 'setup_complete')):
		return

	config = get_notification_config()

	groups = list(config.get("for_doctype")) + list(config.get("for_module"))
	cache = frappe.cache()

	notification_count = {}
	notification_percent = {}

	for name in groups:
		count = cache.hget("notification_count:" + name, frappe.session.user)
		if count is not None:
			notification_count[name] = count

	return {
		"open_count_doctype": get_notifications_for_doctypes(config, notification_count),
		"open_count_module": get_notifications_for_modules(config, notification_count),
		"open_count_other": get_notifications_for_other(config, notification_count),
		"targets": get_notifications_for_targets(config, notification_percent),
		"new_messages": get_new_messages()
	}

def get_new_messages():
	last_update = frappe.cache().hget("notifications_last_update", frappe.session.user)
	now_timestamp = now()
	frappe.cache().hset("notifications_last_update", frappe.session.user, now_timestamp)

	if not last_update:
		return []

	if last_update and time_diff_in_seconds(now_timestamp, last_update) > 1800:
		# no update for 30 mins, consider only the last 30 mins
		last_update = (now_datetime() - relativedelta(seconds=1800)).strftime(DATETIME_FORMAT)

	return frappe.db.sql("""select sender_full_name, content
		from `tabCommunication`
			where communication_type in ('Chat', 'Notification')
			and reference_doctype='user'
			and reference_name = %s
			and creation > %s
			order by creation desc""", (frappe.session.user, last_update), as_dict=1)

def get_notifications_for_modules(config, notification_count):
	"""Notifications for modules"""
	return get_notifications_for("for_module", config, notification_count)

def get_notifications_for_other(config, notification_count):
	"""Notifications for other items"""
	return get_notifications_for("for_other", config, notification_count)

def get_notifications_for(notification_type, config, notification_count):
	open_count = {}
	notification_map = config.get(notification_type) or {}
	for m in notification_map:
		try:
			if m in notification_count:
				open_count[m] = notification_count[m]
			else:
				open_count[m] = frappe.get_attr(notification_map[m])()

				frappe.cache().hset("notification_count:" + m, frappe.session.user, open_count[m])
		except frappe.PermissionError:
			frappe.clear_messages()
			pass
			# frappe.msgprint("Permission Error in notifications for {0}".format(m))

	return open_count

def get_notifications_for_doctypes(config, notification_count):
	"""Notifications for DocTypes"""
	can_read = frappe.get_user().get_can_read()
	open_count_doctype = {}

	for d in config.for_doctype:
		if d in can_read:
			condition = config.for_doctype[d]

			if d in notification_count:
				open_count_doctype[d] = notification_count[d]
			else:
				try:
					if isinstance(condition, dict):
						result = len(frappe.get_list(d, fields=["name"],
							filters=condition, limit_page_length = 100, as_list=True, ignore_ifnull=True))
					else:
						result = frappe.get_attr(condition)()

				except frappe.PermissionError:
					frappe.clear_messages()
					pass
					# frappe.msgprint("Permission Error in notifications for {0}".format(d))

				except Exception as e:
					# OperationalError: (1412, 'Table definition has changed, please retry transaction')
					if e.args[0]!=1412:
						raise

				else:
					open_count_doctype[d] = result
					frappe.cache().hset("notification_count:" + d, frappe.session.user, result)

	return open_count_doctype

def get_notifications_for_targets(config, notification_percent):
	"""Notifications for doc targets"""
	can_read = frappe.get_user().get_can_read()
	doc_target_percents = {}

	# doc_target_percents = {
	# 	"Company": {
	# 		"Acme": 87,
	# 		"RobotsRUs": 50,
	# 	}, {}...
	# }

	for doctype in config.targets:
		if doctype in can_read:
			if doctype in notification_percent:
				doc_target_percents[doctype] = notification_percent[doctype]
			else:
				doc_target_percents[doctype] = {}
				d = config.targets[doctype]
				condition = d["filters"]
				target_field = d["target_field"]
				value_field = d["value_field"]
				try:
					if isinstance(condition, dict):
						doc_list = frappe.get_list(doctype, fields=["name", target_field, value_field],
							filters=condition, limit_page_length = 100, ignore_ifnull=True)

				except frappe.PermissionError:
					frappe.clear_messages()
					pass
				except Exception as e:
					if e.args[0]!=1412:
						raise

				else:
					for doc in doc_list:
						value = doc[value_field]
						target = doc[target_field]
						doc_target_percents[doctype][doc.name] = (value/target * 100) if value < target else 100

	return doc_target_percents

def clear_notifications(user=None):
	if frappe.flags.in_install:
		return

	config = get_notification_config()
	for_doctype = list(config.get('for_doctype')) if config.get('for_doctype') else []
	for_module = list(config.get('for_module')) if config.get('for_module') else []
	groups = for_doctype + for_module
	cache = frappe.cache()

	for name in groups:
		if user:
			cache.hdel("notification_count:" + name, user)
		else:
			cache.delete_key("notification_count:" + name)

	frappe.publish_realtime('clear_notifications')

def delete_notification_count_for(doctype):
	frappe.cache().delete_key("notification_count:" + doctype)
	frappe.publish_realtime('clear_notifications')

def clear_doctype_notifications(doc, method=None, *args, **kwargs):
	config = get_notification_config()
	if isinstance(doc, string_types):
		doctype = doc # assuming doctype name was passed directly
	else:
		doctype = doc.doctype

	if doctype in config.for_doctype:
		delete_notification_count_for(doctype)
		return

def get_notification_info_for_boot():
	out = get_notifications()
	config = get_notification_config()
	can_read = frappe.get_user().get_can_read()
	conditions = {}
	module_doctypes = {}
	doctype_info = dict(frappe.db.sql("""select name, module from tabDocType"""))

	for d in list(set(can_read + list(config.for_doctype))):
		if d in config.for_doctype:
			conditions[d] = config.for_doctype[d]

		if d in doctype_info:
			module_doctypes.setdefault(doctype_info[d], []).append(d)

	out.update({
		"conditions": conditions,
		"module_doctypes": module_doctypes,
	})

	return out

def get_notification_config():
	def _get():
		config = frappe._dict()
		hooks = frappe.get_hooks()
		if hooks:
			for notification_config in hooks.notification_config:
				nc = frappe.get_attr(notification_config)()
				for key in ("for_doctype", "for_module", "for_other", "targets"):
					config.setdefault(key, {})
					config[key].update(nc.get(key, {}))
		return config

	return frappe.cache().get_value("notification_config", _get)

def get_filters_for(doctype):
	'''get open filters for doctype'''
	config = get_notification_config()
	return config.get('for_doctype').get(doctype, {})

@frappe.whitelist()
def get_open_count(doctype, name, items=[]):
	'''Get open count for given transactions and filters

	:param doctype: Reference DocType
	:param name: Reference Name
	:param transactions: List of transactions (json/dict)
	:param filters: optional filters (json/list)'''

	frappe.has_permission(doc=frappe.get_doc(doctype, name), throw=True)

	meta = frappe.get_meta(doctype)
	links = meta.get_dashboard_data()

	# compile all items in a list
	if not items:
		for group in links.transactions:
			items.extend(group.get('items'))

	if not isinstance(items, list):
		items = json.loads(items)

	out = []
	for d in items:
		if d in links.get('internal_links', {}):
			# internal link
			continue

		filters = get_filters_for(d)
		fieldname = links.get('non_standard_fieldnames', {}).get(d, links.fieldname)
		data = {'name': d}
		if filters:
			# get the fieldname for the current document
			# we only need open documents related to the current document
			filters[fieldname] = name
			total = len(frappe.get_all(d, fields='name',
				filters=filters, limit=100, distinct=True, ignore_ifnull=True))
			data['open_count'] = total

		total = len(frappe.get_all(d, fields='name',
			filters={fieldname: name}, limit=100, distinct=True, ignore_ifnull=True))
		data['count'] = total
		out.append(data)

	out = {
		'count': out,
	}

	module = frappe.get_meta_module(doctype)
	if hasattr(module, 'get_timeline_data'):
		out['timeline_data'] = module.get_timeline_data(doctype, name)

	return out
<commit_msg>fix: return empty dict if site is in installation state or setup wiz not completed<commit_after># Copyright (c) 2015, Frappe Technologies Pvt. Ltd. and Contributors
# MIT License. See license.txt

from __future__ import unicode_literals

import frappe
from frappe.utils import time_diff_in_seconds, now, now_datetime, DATETIME_FORMAT
from dateutil.relativedelta import relativedelta
from six import string_types
import json

@frappe.whitelist()
def get_notifications():
	if (frappe.flags.in_install or
		not frappe.db.get_single_value('System Settings', 'setup_complete')):
		return { }

	config = get_notification_config()

	groups = list(config.get("for_doctype")) + list(config.get("for_module"))
	cache = frappe.cache()

	notification_count = {}
	notification_percent = {}

	for name in groups:
		count = cache.hget("notification_count:" + name, frappe.session.user)
		if count is not None:
			notification_count[name] = count

	return {
		"open_count_doctype": get_notifications_for_doctypes(config, notification_count),
		"open_count_module": get_notifications_for_modules(config, notification_count),
		"open_count_other": get_notifications_for_other(config, notification_count),
		"targets": get_notifications_for_targets(config, notification_percent),
		"new_messages": get_new_messages()
	}

def get_new_messages():
	last_update = frappe.cache().hget("notifications_last_update", frappe.session.user)
	now_timestamp = now()
	frappe.cache().hset("notifications_last_update", frappe.session.user, now_timestamp)

	if not last_update:
		return []

	if last_update and time_diff_in_seconds(now_timestamp, last_update) > 1800:
		# no update for 30 mins, consider only the last 30 mins
		last_update = (now_datetime() - relativedelta(seconds=1800)).strftime(DATETIME_FORMAT)

	return frappe.db.sql("""select sender_full_name, content
		from `tabCommunication`
			where communication_type in ('Chat', 'Notification')
			and reference_doctype='user'
			and reference_name = %s
			and creation > %s
			order by creation desc""", (frappe.session.user, last_update), as_dict=1)

def get_notifications_for_modules(config, notification_count):
	"""Notifications for modules"""
	return get_notifications_for("for_module", config, notification_count)

def get_notifications_for_other(config, notification_count):
	"""Notifications for other items"""
	return get_notifications_for("for_other", config, notification_count)

def get_notifications_for(notification_type, config, notification_count):
	open_count = {}
	notification_map = config.get(notification_type) or {}
	for m in notification_map:
		try:
			if m in notification_count:
				open_count[m] = notification_count[m]
			else:
				open_count[m] = frappe.get_attr(notification_map[m])()

				frappe.cache().hset("notification_count:" + m, frappe.session.user, open_count[m])
		except frappe.PermissionError:
			frappe.clear_messages()
			pass
			# frappe.msgprint("Permission Error in notifications for {0}".format(m))

	return open_count

def get_notifications_for_doctypes(config, notification_count):
	"""Notifications for DocTypes"""
	can_read = frappe.get_user().get_can_read()
	open_count_doctype = {}

	for d in config.for_doctype:
		if d in can_read:
			condition = config.for_doctype[d]

			if d in notification_count:
				open_count_doctype[d] = notification_count[d]
			else:
				try:
					if isinstance(condition, dict):
						result = len(frappe.get_list(d, fields=["name"],
							filters=condition, limit_page_length = 100, as_list=True, ignore_ifnull=True))
					else:
						result = frappe.get_attr(condition)()

				except frappe.PermissionError:
					frappe.clear_messages()
					pass
					# frappe.msgprint("Permission Error in notifications for {0}".format(d))

				except Exception as e:
					# OperationalError: (1412, 'Table definition has changed, please retry transaction')
					if e.args[0]!=1412:
						raise

				else:
					open_count_doctype[d] = result
					frappe.cache().hset("notification_count:" + d, frappe.session.user, result)

	return open_count_doctype

def get_notifications_for_targets(config, notification_percent):
	"""Notifications for doc targets"""
	can_read = frappe.get_user().get_can_read()
	doc_target_percents = {}

	# doc_target_percents = {
	# 	"Company": {
	# 		"Acme": 87,
	# 		"RobotsRUs": 50,
	# 	}, {}...
	# }

	for doctype in config.targets:
		if doctype in can_read:
			if doctype in notification_percent:
				doc_target_percents[doctype] = notification_percent[doctype]
			else:
				doc_target_percents[doctype] = {}
				d = config.targets[doctype]
				condition = d["filters"]
				target_field = d["target_field"]
				value_field = d["value_field"]
				try:
					if isinstance(condition, dict):
						doc_list = frappe.get_list(doctype, fields=["name", target_field, value_field],
							filters=condition, limit_page_length = 100, ignore_ifnull=True)

				except frappe.PermissionError:
					frappe.clear_messages()
					pass
				except Exception as e:
					if e.args[0]!=1412:
						raise

				else:
					for doc in doc_list:
						value = doc[value_field]
						target = doc[target_field]
						doc_target_percents[doctype][doc.name] = (value/target * 100) if value < target else 100

	return doc_target_percents

def clear_notifications(user=None):
	if frappe.flags.in_install:
		return

	config = get_notification_config()
	for_doctype = list(config.get('for_doctype')) if config.get('for_doctype') else []
	for_module = list(config.get('for_module')) if config.get('for_module') else []
	groups = for_doctype + for_module
	cache = frappe.cache()

	for name in groups:
		if user:
			cache.hdel("notification_count:" + name, user)
		else:
			cache.delete_key("notification_count:" + name)

	frappe.publish_realtime('clear_notifications')

def delete_notification_count_for(doctype):
	frappe.cache().delete_key("notification_count:" + doctype)
	frappe.publish_realtime('clear_notifications')

def clear_doctype_notifications(doc, method=None, *args, **kwargs):
	config = get_notification_config()
	if isinstance(doc, string_types):
		doctype = doc # assuming doctype name was passed directly
	else:
		doctype = doc.doctype

	if doctype in config.for_doctype:
		delete_notification_count_for(doctype)
		return

def get_notification_info_for_boot():
	out = get_notifications()
	config = get_notification_config()
	can_read = frappe.get_user().get_can_read()
	conditions = {}
	module_doctypes = {}
	doctype_info = dict(frappe.db.sql("""select name, module from tabDocType"""))

	for d in list(set(can_read + list(config.for_doctype))):
		if d in config.for_doctype:
			conditions[d] = config.for_doctype[d]

		if d in doctype_info:
			module_doctypes.setdefault(doctype_info[d], []).append(d)

	out.update({
		"conditions": conditions,
		"module_doctypes": module_doctypes,
	})

	return out

def get_notification_config():
	def _get():
		config = frappe._dict()
		hooks = frappe.get_hooks()
		if hooks:
			for notification_config in hooks.notification_config:
				nc = frappe.get_attr(notification_config)()
				for key in ("for_doctype", "for_module", "for_other", "targets"):
					config.setdefault(key, {})
					config[key].update(nc.get(key, {}))
		return config

	return frappe.cache().get_value("notification_config", _get)

def get_filters_for(doctype):
	'''get open filters for doctype'''
	config = get_notification_config()
	return config.get('for_doctype').get(doctype, {})

@frappe.whitelist()
def get_open_count(doctype, name, items=[]):
	'''Get open count for given transactions and filters

	:param doctype: Reference DocType
	:param name: Reference Name
	:param transactions: List of transactions (json/dict)
	:param filters: optional filters (json/list)'''

	frappe.has_permission(doc=frappe.get_doc(doctype, name), throw=True)

	meta = frappe.get_meta(doctype)
	links = meta.get_dashboard_data()

	# compile all items in a list
	if not items:
		for group in links.transactions:
			items.extend(group.get('items'))

	if not isinstance(items, list):
		items = json.loads(items)

	out = []
	for d in items:
		if d in links.get('internal_links', {}):
			# internal link
			continue

		filters = get_filters_for(d)
		fieldname = links.get('non_standard_fieldnames', {}).get(d, links.fieldname)
		data = {'name': d}
		if filters:
			# get the fieldname for the current document
			# we only need open documents related to the current document
			filters[fieldname] = name
			total = len(frappe.get_all(d, fields='name',
				filters=filters, limit=100, distinct=True, ignore_ifnull=True))
			data['open_count'] = total

		total = len(frappe.get_all(d, fields='name',
			filters={fieldname: name}, limit=100, distinct=True, ignore_ifnull=True))
		data['count'] = total
		out.append(data)

	out = {
		'count': out,
	}

	module = frappe.get_meta_module(doctype)
	if hasattr(module, 'get_timeline_data'):
		out['timeline_data'] = module.get_timeline_data(doctype, name)

	return out
<|endoftext|>


<commit_before>import collections

from puzzle.heuristics import analyze_word
from puzzle.problems import problem

class AnagramProblem(problem.Problem):

  @staticmethod
  def score(lines):
    # TODO: Support more input.
    if len(lines) > 1:
      return 0
    return analyze_word.score_anagram(lines[0])

  def _solve(self):
    index = analyze_word.get_anagram_index()
    results = collections.OrderedDict()
    if self.lines[0] not in index:
      return results
    for index, word in reversed(list(enumerate(index[self.lines[0]]))):
      # TODO: Use frequency for score.
      results[word] = index + 1
    return results
<commit_msg>Use warehouse in anagram_problem.py<commit_after>import collections

from data import warehouse
from puzzle.heuristics import analyze_word
from puzzle.problems import problem

class AnagramProblem(problem.Problem):

  @staticmethod
  def score(lines):
    # TODO: Support more input.
    if len(lines) > 1:
      return 0
    return analyze_word.score_anagram(lines[0])

  def _solve(self):
    index = warehouse.get('/words/unigram/anagram_index')
    results = collections.OrderedDict()
    if self.lines[0] not in index:
      return results
    for index, word in reversed(list(enumerate(index[self.lines[0]]))):
      # TODO: Use frequency for score.
      results[word] = index + 1
    return results
<|endoftext|>


<commit_before># -*- coding: UTF-8 -*-
# Copyright (C) 2009 Itaapy, ArsAperta, Pierlis, Talend

# Import from the Standard Library
from unittest import TestCase, main

# Import from lpod
from lpod.document import odf_get_document
from lpod.style import odf_create_style
from lpod.styles import hex2rgb, rgb2hex

class Hex2RgbTestCase(TestCase):

    def test_color_low(self):
        color = '#012345'
        expected = (1, 35, 69)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_high(self):
        color = '#ABCDEF'
        expected = (171, 205, 239)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_lowercase(self):
        color = '#abcdef'
        expected = (171, 205, 239)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_bad_size(self):
        color = '#fff'
        self.assertRaises(ValueError, hex2rgb, color)

    def test_color_bad_format(self):
        color = '978EAE'
        self.assertRaises(ValueError, hex2rgb, color)

    def test_color_bad_hex(self):
        color = '#978EAZ'
        self.assertRaises(ValueError, hex2rgb, color)

class Rgb2HexTestCase(TestCase):

    def test_color_name(self):
        color = 'violet'
        expected = '#EE82EE'
        self.assertEqual(rgb2hex(color), expected)

    def test_color_tuple(self):
        color = (171, 205, 239)
        expected = '#ABCDEF'
        self.assertEqual(rgb2hex(color), expected)

    def test_color_bad_name(self):
        color = 'dark white'
        self.assertRaises(KeyError, rgb2hex, color)

    def test_color_bad_tuple(self):
        # For alpha channel? ;-)
        color = (171, 205, 238, 128)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_low_channel(self):
        color = (171, 205, -1)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_high_channel(self):
        color = (171, 205, 256)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_value(self):
        color = {}
        self.assertRaises(TypeError, rgb2hex, color)

class TestStyle(TestCase):

    def setUp(self):
        self.document = document = odf_get_document('samples/example.odt')
        self.styles = document.get_xmlpart('styles')

    def tearDown(self):
        del self.styles
        del self.document

    def test_create_style(self):
        style = odf_create_style('paragraph', 'style1')
        expected = ('<style:style style:name="style1" '
                      'style:family="paragraph"/>')
        self.assertEqual(style.serialize(), expected)

    def test_get_style_list(self):
        style_list = self.styles.get_style_list()
        self.assertEqual(len(style_list), 15)

    def test_get_style_list_paragraph(self):
        style_list = self.styles.get_style_list(family='paragraph')
        self.assertEqual(len(style_list), 10)

    def test_get_style_list_master_page(self):
        style_list = self.styles.get_style_list(family='master-page')
        self.assertEqual(len(style_list), 1)

    def test_get_style_automatic(self):
        style = self.styles.get_style('page-layout', u'Mpm1')
        self.assertNotEqual(style, None)

    def test_get_style_named(self):
        style = self.styles.get_style('paragraph', u'Heading_20_1')
        self.assertEqual(style.get_style_display_name(), u"Heading 1")

    def test_get_style_display_name(self):
        style = self.styles.get_style('paragraph', u"Text body",
                                      display_name=True)
        self.assertEqual(style.get_style_name(), u"Text_20_body")

    def test_insert_style(self):
        styles = self.styles.clone()
        style = odf_create_style('paragraph', u'style1', area='text',
                                 **{'fo:color': '#0000ff',
                                    'fo:background-color': '#ff0000'})
        context = styles.get_element('//office:styles')
        context.append_element(style)

        expected = ('<style:style style:name="style1" '
                                  'style:family="paragraph">\n'
                    '  <style:text-properties fo:color="#0000ff" '
                                             'fo:background-color="#ff0000"/>\n'
                    '</style:style>\n')
        get1 = styles.get_style('paragraph', u'style1')
        self.assertEqual(get1.serialize(pretty=True), expected)

class TestInsertStyleCase(TestCase):

    def setUp(self):
        self.doc = odf_get_document('samples/example.odt')

    def test_insert_common_style(self):
        doc = self.doc

        style = odf_create_style('paragraph', u'MyStyle')
        doc.insert_style(style)
        inserted_style = doc.get_style('paragraph', u'MyStyle')

        self.assertEqual(style.serialize(), inserted_style.serialize())

    def test_insert_default_style(self):
        doc = self.doc

        style = odf_create_style('paragraph', u'MyStyle')
        #doc.insert_style(style, default=True)

        inserted_style = doc.get_style('paragraph')
        expected = '<style:default-style style:family="paragraph"/>'

        self.assertEqual(inserted_style.serialize(), expected)

if __name__ == '__main__':
    main()
<commit_msg>python: test: styles: Add some tests about "document.insert_style"<commit_after># -*- coding: UTF-8 -*-
# Copyright (C) 2009 Itaapy, ArsAperta, Pierlis, Talend

# Import from the Standard Library
from unittest import TestCase, main

# Import from lpod
from lpod.document import odf_get_document
from lpod.style import odf_create_style
from lpod.styles import hex2rgb, rgb2hex

class Hex2RgbTestCase(TestCase):

    def test_color_low(self):
        color = '#012345'
        expected = (1, 35, 69)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_high(self):
        color = '#ABCDEF'
        expected = (171, 205, 239)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_lowercase(self):
        color = '#abcdef'
        expected = (171, 205, 239)
        self.assertEqual(hex2rgb(color), expected)

    def test_color_bad_size(self):
        color = '#fff'
        self.assertRaises(ValueError, hex2rgb, color)

    def test_color_bad_format(self):
        color = '978EAE'
        self.assertRaises(ValueError, hex2rgb, color)

    def test_color_bad_hex(self):
        color = '#978EAZ'
        self.assertRaises(ValueError, hex2rgb, color)

class Rgb2HexTestCase(TestCase):

    def test_color_name(self):
        color = 'violet'
        expected = '#EE82EE'
        self.assertEqual(rgb2hex(color), expected)

    def test_color_tuple(self):
        color = (171, 205, 239)
        expected = '#ABCDEF'
        self.assertEqual(rgb2hex(color), expected)

    def test_color_bad_name(self):
        color = 'dark white'
        self.assertRaises(KeyError, rgb2hex, color)

    def test_color_bad_tuple(self):
        # For alpha channel? ;-)
        color = (171, 205, 238, 128)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_low_channel(self):
        color = (171, 205, -1)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_high_channel(self):
        color = (171, 205, 256)
        self.assertRaises(ValueError, rgb2hex, color)

    def test_color_bad_value(self):
        color = {}
        self.assertRaises(TypeError, rgb2hex, color)

class TestStyle(TestCase):

    def setUp(self):
        self.document = document = odf_get_document('samples/example.odt')
        self.styles = document.get_xmlpart('styles')

    def tearDown(self):
        del self.styles
        del self.document

    def test_create_style(self):
        style = odf_create_style('paragraph', 'style1')
        expected = ('<style:style style:name="style1" '
                      'style:family="paragraph"/>')
        self.assertEqual(style.serialize(), expected)

    def test_get_style_list(self):
        style_list = self.styles.get_style_list()
        self.assertEqual(len(style_list), 15)

    def test_get_style_list_paragraph(self):
        style_list = self.styles.get_style_list(family='paragraph')
        self.assertEqual(len(style_list), 10)

    def test_get_style_list_master_page(self):
        style_list = self.styles.get_style_list(family='master-page')
        self.assertEqual(len(style_list), 1)

    def test_get_style_automatic(self):
        style = self.styles.get_style('page-layout', u'Mpm1')
        self.assertNotEqual(style, None)

    def test_get_style_named(self):
        style = self.styles.get_style('paragraph', u'Heading_20_1')
        self.assertEqual(style.get_style_display_name(), u"Heading 1")

    def test_get_style_display_name(self):
        style = self.styles.get_style('paragraph', u"Text body",
                                      display_name=True)
        self.assertEqual(style.get_style_name(), u"Text_20_body")

    def test_insert_style(self):
        styles = self.styles.clone()
        style = odf_create_style('paragraph', u'style1', area='text',
                                 **{'fo:color': '#0000ff',
                                    'fo:background-color': '#ff0000'})
        context = styles.get_element('//office:styles')
        context.append_element(style)

        expected = ('<style:style style:name="style1" '
                                  'style:family="paragraph">\n'
                    '  <style:text-properties fo:color="#0000ff" '
                                             'fo:background-color="#ff0000"/>\n'
                    '</style:style>\n')
        get1 = styles.get_style('paragraph', u'style1')
        self.assertEqual(get1.serialize(pretty=True), expected)

class TestInsertStyleCase(TestCase):

    def setUp(self):
        self.doc = odf_get_document('samples/example.odt')

    def test_insert_common_style(self):
        doc = self.doc

        style = odf_create_style('paragraph', u'MyStyle')
        doc.insert_style(style)
        inserted_style = doc.get_style('paragraph', u'MyStyle')

        self.assertEqual(style.serialize(), inserted_style.serialize())

    def test_insert_default_style(self):
        doc = self.doc

        style = odf_create_style('paragraph', u'MyStyle')
        doc.insert_style(style, default=True)

        inserted_style = doc.get_style('paragraph')
        expected = '<style:default-style style:family="paragraph"/>'

        self.assertEqual(inserted_style.serialize(), expected)

    def test_insert_automatic_style(self):
        doc = self.doc

        style = odf_create_style('paragraph')
        doc.insert_style(style, automatic=True)
        self.assertNotEqual(style.get_style_name(), None)

    def test_insert_with_error(self):
        doc = self.doc

        style = odf_create_style('paragraph', u'MyStyle')
        self.assertRaises(AttributeError, doc.insert_style,
                          style=style, automatic=True, default=True)

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


<commit_before># -* encoding: utf-8 *-

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import configargparse
import tempfile
import os

from typing import Any, Type, List

from gopythongo.builders import BaseBuilder, get_dependencies
from gopythongo.utils import print_info, highlight, run_process, print_debug, ErrorMessage, cmdargs_unquote_split
from gopythongo.utils.buildcontext import the_context

class PbuilderBuilder(BaseBuilder):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

    @property
    def builder_name(self) -> str:
        return "pbuilder"

    def add_args(self, parser: configargparse.ArgumentParser) -> None:
        gr_pbuilder = parser.add_argument_group("Pbuilder Builder options")
        gr_pbuilder.add_argument("--use-pbuilder", dest="pbuilder_executable", default="/usr/sbin/pbuilder",
                                 env_var="PBUILDER_EXECUTABLE",
                                 help="Specify an alternative pbuilder executable")
        gr_pbuilder.add_argument("--basetgz", dest="basetgz", default="/var/cache/pbuilder/base.tgz", env_var="BASETGZ",
                                 help="Cache and reuse the pbuilder base environment. gopythongo will call pbuilder "
                                      "create on this file if it doesn't exist")
        gr_pbuilder.add_argument("--distribution", dest="pbuilder_distribution", default=None, env_var="DISTRIBUTION",
                                 help="Use this distribution for creating the pbuilder environment using debootstrap.")
        gr_pbuilder.add_argument("--pbuilder-force-recreate", dest="pbuilder_force_recreate", action="store_true",
                                 help="Delete the base environment if it exists already")
        gr_pbuilder.add_argument("--pbuilder-reprovision", dest="pbuilder_reprovision", action="store_true",
                                 default=False,
                                 help="Run all --run-after-create commands regardless of whether the pbuilder base "
                                      "environment already exists.")
        gr_pbuilder.add_argument("--pbuilder-opts", dest="pbuilder_opts", default="", env_var="PBUILDER_OPTS",
                                 help="Options which will be put into every pbuilder command-line executed by "
                                      "GoPythonGo")
        gr_pbuilder.add_argument("--pbuilder-create-opts", dest="pbuilder_create_opts", default="",
                                 env_var="PBUILDER_CREATE_OPTS",
                                 help="Options which will be appended to the pbuilder --create command-line")
        gr_pbuilder.add_argument("--pbuilder-execute-opts", dest="pbuilder_execute_opts", default="",
                                 env_var="PBUILDER_EXECUTE_OPTS",
                                 help="Options which will be appended to the pbuilder --execute command-line")
        gr_pbuilder.add_argument("--apt-get", dest="build_deps", action="append", default=[],
                                 help="Packages to install using apt-get prior to creating the virtualenv (e.g. driver "
                                      "libs for databases so that Python C extensions compile correctly")
        gr_pbuilder.add_argument("--no-install-defaults", dest="install_defaults", action="store_false",
                                 default=True,
                                 help="By default GoPythonGo will always install python, python-virtualenv, "
                                      "python-pip, python[3]-dev, virtualenv and possibly eatmydata. If you set this "
                                      "flag you will have to install python using --apt-get, or GoPythonGo will not be "
                                      "able to run inside the container, but this gives you more control about what "
                                      "Python version runs.")

    def validate_args(self, args: configargparse.Namespace) -> None:
        if args.is_inner:
            pass
        else:
            # validate things we only need to validate in the outer layer
            if not os.path.exists(args.pbuilder_executable) or not os.access(args.pbuilder_executable, os.X_OK):
                raise ErrorMessage("pbuilder not found in path or not executable (%s).\n"
                                   "You can specify an alternative path using %s" %
                                   (args.pbuilder_executable, highlight("--use-pbuilder")))

            if args.basetgz and os.path.exists(args.basetgz) and not os.path.isfile(args.basetgz):
                raise ErrorMessage("pbuilder basetgz %s exists but is not a file. Can't continue with this "
                                   "inconsistency." % highlight(args.basetgz))

            if not args.pbuilder_distribution:
                raise ErrorMessage("pbuilder distribution unfortunately defaults to %s, so you must explicitly set it "
                                   "using the parameter %s" %
                                   (highlight("sid (unstable)"), highlight("--distribution")))
            elif ("debian/%s" % args.pbuilder_distribution) not in get_dependencies() \
                    and args.install_defaults:
                raise ErrorMessage("GoPythonGo does not have stored package dependencies for distribution %s. You can "
                                   "set %s to ignore this error." %
                                   (highlight(args.pbuilder_distribution), highlight("--no-install-defaults")))

            if os.getuid() != 0:
                raise ErrorMessage("pbuilder requires root privileges. Please run GoPythonGo as root when using "
                                   "pbuilder")

    def build(self, args: configargparse.Namespace) -> None:
        print_info("Building with %s" % highlight("pbuilder"))

        do_create = True
        if args.basetgz and os.path.exists(args.basetgz) and not args.pbuilder_force_recreate:
            do_create = False

        if args.basetgz and os.path.exists(args.basetgz) and args.pbuilder_force_recreate:
            os.unlink(args.basetgz)

        if do_create:
            create_cmdline = [args.pbuilder_executable, "--create"]
            create_cmdline += cmdargs_unquote_split(args.pbuilder_opts)
            create_cmdline += cmdargs_unquote_split(args.pbuilder_create_opts)
            if args.pbuilder_distribution:
                create_cmdline += ["--distribution", args.pbuilder_distribution]

            if args.basetgz:
                create_cmdline += ["--basetgz", args.basetgz]

            if args.install_defaults:
                args.build_deps += get_dependencies()["debian/%s" % args.pbuilder_distribution]
                if args.eatmydata:
                    args.build_deps += ["eatmydata"]

            if args.build_deps:
                create_cmdline += ["--extrapackages", " ".join(args.build_deps)]

            run_process(*create_cmdline)

        build_args = []  # type: List[str]
        build_args += cmdargs_unquote_split(args.pbuilder_opts)
        build_args += cmdargs_unquote_split(args.pbuilder_execute_opts)

        for mount in args.mounts + list(the_context.mounts):
            build_args += ["--bindmounts", mount]

        if args.basetgz:
            build_args += ["--basetgz", args.basetgz]

        if do_create or args.pbuilder_reprovision:
            for ix, runspec in enumerate(args.run_after_create):
                print_info("Running post-creation commands for build environment %s of %s" %
                           (highlight(str(ix + 1)), highlight(str(len(args.run_after_create)))))
                if os.path.isfile(os.path.abspath(runspec)):
                    runspec = os.path.abspath(runspec)
                post_create_cmdline = [args.pbuilder_executable, "--execute"] + build_args + \
                                      ["--save-after-exec", "--", runspec]
                run_process(*post_create_cmdline)

        if args.builder_debug_login:
            build_cmdline = [args.pbuilder_executable, "--login"] + build_args
            debug_cmdline = build_cmdline + ["--"] + the_context.get_gopythongo_inner_commandline()
            debug_cmdline = [x if x != "--login" else "--execute" for x in debug_cmdline]
            print_debug("Without --builder-debug-login, GoPythonGo would have run: %s" % " ".join(debug_cmdline))
        else:
            # pbuilder COPIES the first argument after "--" to the chroot and executes it, so we can't reference
            # the GoPythonGo python interpreter directly. Instead we need to create a intermediary script file
            from gopythongo.main import tempfiles
            build_cmdline = [args.pbuilder_executable, "--execute"] + build_args + ["--"]
            scriptfd, scriptfn = tempfile.mkstemp()
            tempfiles.append(scriptfn)
            with open(scriptfd, "wt", encoding="utf-8") as f:
                print("#!/bin/sh", file=f)
                print(" ".join(the_context.get_gopythongo_inner_commandline()), file=f)
            print_debug("Running the following command inside the build environment: %s" %
                        " ".join(the_context.get_gopythongo_inner_commandline()))
            build_cmdline += [scriptfn]

        run_process(*build_cmdline, interactive=args.builder_debug_login)

    def print_help(self) -> None:
        print("Pbuilder Builder\n"
              "================\n"
              "\n"
              "Builds virtualenvs in a chroot using Debian's pbuilder. This has the drawback,\n"
              "that GoPythonGo needs to run as root to utilize pbuilder correctly, as it needs\n"
              "chroot privileges (you might be able to do something using fakeroot and setcap\n"
              "cap_sys_chroot, but that's untested).\n")

builder_class = PbuilderBuilder  # type: Type[PbuilderBuilder]
<commit_msg>try our best to clean up when we created an unfinished pbuilder environment<commit_after># -* encoding: utf-8 *-

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import configargparse
import tempfile
import os

from typing import Any, Type, List

from gopythongo.builders import BaseBuilder, get_dependencies
from gopythongo.utils import print_info, highlight, run_process, print_debug, ErrorMessage, cmdargs_unquote_split
from gopythongo.utils.buildcontext import the_context

class PbuilderBuilder(BaseBuilder):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

    @property
    def builder_name(self) -> str:
        return "pbuilder"

    def add_args(self, parser: configargparse.ArgumentParser) -> None:
        gr_pbuilder = parser.add_argument_group("Pbuilder Builder options")
        gr_pbuilder.add_argument("--use-pbuilder", dest="pbuilder_executable", default="/usr/sbin/pbuilder",
                                 env_var="PBUILDER_EXECUTABLE",
                                 help="Specify an alternative pbuilder executable")
        gr_pbuilder.add_argument("--basetgz", dest="basetgz", default="/var/cache/pbuilder/base.tgz", env_var="BASETGZ",
                                 help="Cache and reuse the pbuilder base environment. gopythongo will call pbuilder "
                                      "create on this file if it doesn't exist")
        gr_pbuilder.add_argument("--distribution", dest="pbuilder_distribution", default=None, env_var="DISTRIBUTION",
                                 help="Use this distribution for creating the pbuilder environment using debootstrap.")
        gr_pbuilder.add_argument("--pbuilder-force-recreate", dest="pbuilder_force_recreate", action="store_true",
                                 help="Delete the base environment if it exists already")
        gr_pbuilder.add_argument("--pbuilder-reprovision", dest="pbuilder_reprovision", action="store_true",
                                 default=False,
                                 help="Run all --run-after-create commands regardless of whether the pbuilder base "
                                      "environment already exists.")
        gr_pbuilder.add_argument("--pbuilder-opts", dest="pbuilder_opts", default="", env_var="PBUILDER_OPTS",
                                 help="Options which will be put into every pbuilder command-line executed by "
                                      "GoPythonGo")
        gr_pbuilder.add_argument("--pbuilder-create-opts", dest="pbuilder_create_opts", default="",
                                 env_var="PBUILDER_CREATE_OPTS",
                                 help="Options which will be appended to the pbuilder --create command-line")
        gr_pbuilder.add_argument("--pbuilder-execute-opts", dest="pbuilder_execute_opts", default="",
                                 env_var="PBUILDER_EXECUTE_OPTS",
                                 help="Options which will be appended to the pbuilder --execute command-line")
        gr_pbuilder.add_argument("--apt-get", dest="build_deps", action="append", default=[],
                                 help="Packages to install using apt-get prior to creating the virtualenv (e.g. driver "
                                      "libs for databases so that Python C extensions compile correctly")
        gr_pbuilder.add_argument("--no-install-defaults", dest="install_defaults", action="store_false",
                                 default=True,
                                 help="By default GoPythonGo will always install python, python-virtualenv, "
                                      "python-pip, python[3]-dev, virtualenv and possibly eatmydata. If you set this "
                                      "flag you will have to install python using --apt-get, or GoPythonGo will not be "
                                      "able to run inside the container, but this gives you more control about what "
                                      "Python version runs.")

    def validate_args(self, args: configargparse.Namespace) -> None:
        if args.is_inner:
            pass
        else:
            # validate things we only need to validate in the outer layer
            if not os.path.exists(args.pbuilder_executable) or not os.access(args.pbuilder_executable, os.X_OK):
                raise ErrorMessage("pbuilder not found in path or not executable (%s).\n"
                                   "You can specify an alternative path using %s" %
                                   (args.pbuilder_executable, highlight("--use-pbuilder")))

            if args.basetgz and os.path.exists(args.basetgz) and not os.path.isfile(args.basetgz):
                raise ErrorMessage("pbuilder basetgz %s exists but is not a file. Can't continue with this "
                                   "inconsistency." % highlight(args.basetgz))

            if not args.pbuilder_distribution:
                raise ErrorMessage("pbuilder distribution unfortunately defaults to %s, so you must explicitly set it "
                                   "using the parameter %s" %
                                   (highlight("sid (unstable)"), highlight("--distribution")))
            elif ("debian/%s" % args.pbuilder_distribution) not in get_dependencies() \
                    and args.install_defaults:
                raise ErrorMessage("GoPythonGo does not have stored package dependencies for distribution %s. You can "
                                   "set %s to ignore this error." %
                                   (highlight(args.pbuilder_distribution), highlight("--no-install-defaults")))

            if os.getuid() != 0:
                raise ErrorMessage("pbuilder requires root privileges. Please run GoPythonGo as root when using "
                                   "pbuilder")

    def build(self, args: configargparse.Namespace) -> None:
        print_info("Building with %s" % highlight("pbuilder"))

        do_create = True
        if args.basetgz and os.path.exists(args.basetgz) and not args.pbuilder_force_recreate:
            do_create = False

        if args.basetgz and os.path.exists(args.basetgz) and args.pbuilder_force_recreate:
            os.unlink(args.basetgz)

        if do_create:
            def pbuilder_cleanup() -> None:
                print_info("Removing probably broken basetgz %s" % args.basetgz)
                os.unlink(args.basetgz)

            import gopythongo.main
            # if anything happens while we provision the basetgz, pbuilder likes to leave the
            # unfinished result lying around.
            gopythongo.main.break_handlers["pbuilder-tgz-cleanup"] = pbuilder_cleanup

            create_cmdline = [args.pbuilder_executable, "--create"]
            create_cmdline += cmdargs_unquote_split(args.pbuilder_opts)
            create_cmdline += cmdargs_unquote_split(args.pbuilder_create_opts)
            if args.pbuilder_distribution:
                create_cmdline += ["--distribution", args.pbuilder_distribution]

            if args.basetgz:
                create_cmdline += ["--basetgz", args.basetgz]

            if args.install_defaults:
                args.build_deps += get_dependencies()["debian/%s" % args.pbuilder_distribution]
                if args.eatmydata:
                    args.build_deps += ["eatmydata"]

            if args.build_deps:
                create_cmdline += ["--extrapackages", " ".join(args.build_deps)]

            run_process(*create_cmdline)

            # once we're here, we're out of the woods, so we remove the cleanup handler
            del gopythongo.main.break_handlers["pbuilder-tgz-cleanup"]

        build_args = []  # type: List[str]
        build_args += cmdargs_unquote_split(args.pbuilder_opts)
        build_args += cmdargs_unquote_split(args.pbuilder_execute_opts)

        for mount in args.mounts + list(the_context.mounts):
            build_args += ["--bindmounts", mount]

        if args.basetgz:
            build_args += ["--basetgz", args.basetgz]

        if do_create or args.pbuilder_reprovision:
            for ix, runspec in enumerate(args.run_after_create):
                print_info("Running post-creation commands for build environment %s of %s" %
                           (highlight(str(ix + 1)), highlight(str(len(args.run_after_create)))))
                if os.path.isfile(os.path.abspath(runspec)):
                    runspec = os.path.abspath(runspec)
                post_create_cmdline = [args.pbuilder_executable, "--execute"] + build_args + \
                                      ["--save-after-exec", "--", runspec]
                run_process(*post_create_cmdline)

        if args.builder_debug_login:
            build_cmdline = [args.pbuilder_executable, "--login"] + build_args
            debug_cmdline = build_cmdline + ["--"] + the_context.get_gopythongo_inner_commandline()
            debug_cmdline = [x if x != "--login" else "--execute" for x in debug_cmdline]
            print_debug("Without --builder-debug-login, GoPythonGo would have run: %s" % " ".join(debug_cmdline))
        else:
            # pbuilder COPIES the first argument after "--" to the chroot and executes it, so we can't reference
            # the GoPythonGo python interpreter directly. Instead we need to create a intermediary script file
            from gopythongo.main import tempfiles
            build_cmdline = [args.pbuilder_executable, "--execute"] + build_args + ["--"]
            scriptfd, scriptfn = tempfile.mkstemp()
            tempfiles.append(scriptfn)
            with open(scriptfd, "wt", encoding="utf-8") as f:
                print("#!/bin/sh", file=f)
                print(" ".join(the_context.get_gopythongo_inner_commandline()), file=f)
            print_debug("Running the following command inside the build environment: %s" %
                        " ".join(the_context.get_gopythongo_inner_commandline()))
            build_cmdline += [scriptfn]

        run_process(*build_cmdline, interactive=args.builder_debug_login)

    def print_help(self) -> None:
        print("Pbuilder Builder\n"
              "================\n"
              "\n"
              "Builds virtualenvs in a chroot using Debian's pbuilder. This has the drawback,\n"
              "that GoPythonGo needs to run as root to utilize pbuilder correctly, as it needs\n"
              "chroot privileges (you might be able to do something using fakeroot and setcap\n"
              "cap_sys_chroot, but that's untested).\n")

builder_class = PbuilderBuilder  # type: Type[PbuilderBuilder]
<|endoftext|>


<commit_before>from __future__ import division
"""
Functions to produce parametric representations of
conic sections in two dimensions. This can be applied
to plotting 2D slices of quadrics as well.

We could generalize this to n dimensions, but that is probably
not super useful at the moment.
"""
import numpy as N
from matplotlib import pyplot as P

defaults = dict(
    center=N.zeros(2),
    n=500)

def hyperbola(axes, **kwargs):
    """
    Plots a hyperbola that opens along y axis
    """
    center = kwargs.pop('center', defaults['center'])
    th = N.linspace(0,2*N.pi,kwargs.pop('n', 500))

    x = axes[0]*N.tan(th)+center[0]
    y = axes[1]*1/N.cos(th)+center[1]

    extrema = [N.argmin(x),N.argmax(x)]

    def remove_asymptotes(arr):
        arr[extrema] = N.nan
        return arr

    xy = tuple(remove_asymptotes(i) for i in (x,y))

    return xy

def ellipse(axes,  **kwargs):
    center = kwargs.pop('center', defaults['center'])
    th = N.linspace(0,2*N.pi,kwargs.pop('n', 500))
    return (
        axes[0]*N.cos(th)+center[0],
        axes[1]*N.sin(th)+center[1])

def __squished_ellipse(axes):
    scalar = 2
    ax1 = axes.copy()
    ax1[0] = scalar*N.sqrt(N.sqrt(3)/2)*(axes[1]**2)/axes[0]
    center = scalar*axes[1]
    return ax1, center

def __inverse_ellipse(axes, scalar=1):
    ax1 = 1/axes*axes[1]**2*scalar
    center = ax1[1]*N.sqrt(2)
    return ax1, center

def __reverse_ellipse(axes, scalar=1):
    """
    This method doesn't work as well
    """
    ax1 = axes.copy()[::-1]*scalar
    center = ax1[1]*N.sqrt(2)*scalar
    return ax1, center

def plot_conjugate_conics(ax, axes, width=None, plot_foci=False):
    hyp, = ax.plot(*hyperbola(axes))
    ax.plot(*ellipse(axes), zorder=-5, color=hyp.get_color(),alpha=0.5, linewidth=1)

    # Normal vector ellipse axes lengths
    # scales inversely with axes but can
    # be rescaled at will

    ax1,center = __squished_ellipse(axes)

    ell, = ax.plot(*ellipse(ax1, center=[0,center]))

    if plot_foci:
        # Plot hyperbola focus
        hyp_c = N.sqrt(N.sum(axes**2))
        ax.plot([0,0],[hyp_c,-hyp_c], '.', color=hyp.get_color())

        # Plot ellipse foci
        c = ax1**2
        c.sort()
        c[0] *= -1
        ell_c = N.sqrt(N.sum(c))
        # Plot ellipse foci
        ax.plot([0,0],[center+ell_c,center-ell_c], '.', color=ell.get_color())

    # Plot tangents
    xvals = N.array([-500,500])
    yvals = axes[1]/axes[0]*xvals
    kw = dict(zorder=-1, color=hyp.get_color(), linewidth=0.5)
    ax.plot(xvals, yvals, ':', **kw)
    ax.plot(xvals, -yvals, ':', **kw)

    #yvals = axes[0]/axes[1]*xvals
    kw = dict(zorder=-1, color=ell.get_color(), linewidth=0.5)
    ax.plot(yvals, xvals, ':', **kw)
    ax.plot(yvals, -xvals, ':', **kw)

    if width is None:
        width = N.linalg.norm(axes)*4
    lim = N.array([-width,width])/2
    ax.set_xlim(lim)
    ax.set_ylim(lim*0.6)
    ax.set_aspect(1)

    return ax
<commit_msg>Better axis limits for parametric plots<commit_after>from __future__ import division
"""
Functions to produce parametric representations of
conic sections in two dimensions. This can be applied
to plotting 2D slices of quadrics as well.

We could generalize this to n dimensions, but that is probably
not super useful at the moment.
"""
import numpy as N
from matplotlib import pyplot as P

defaults = dict(
    center=N.zeros(2),
    n=500)

def hyperbola(axes, **kwargs):
    """
    Plots a hyperbola that opens along y axis
    """
    opens_up = kwargs.pop('opens_up', True)
    center = kwargs.pop('center', defaults['center'])
    th = N.linspace(0,2*N.pi,kwargs.pop('n', 500))

    vals = [N.tan(th),1/N.cos(th)]
    if not opens_up:
        vals = vals[::-1]

    x = axes[0]*vals[0]+center[0]
    y = axes[1]*vals[1]+center[1]

    extrema = [N.argmin(x),N.argmax(x)]

    def remove_asymptotes(arr):
        arr[extrema] = N.nan
        return arr

    xy = tuple(remove_asymptotes(i) for i in (x,y))

    return xy

def ellipse(axes,  **kwargs):
    center = kwargs.pop('center', defaults['center'])
    th = N.linspace(0,2*N.pi,kwargs.pop('n', 500))
    return (
        axes[0]*N.cos(th)+center[0],
        axes[1]*N.sin(th)+center[1])

def __squished_ellipse(axes):
    scalar = 2
    ax1 = axes.copy()
    ax1[0] = scalar*N.sqrt(N.sqrt(3)/2)*(axes[1]**2)/axes[0]
    center = scalar*axes[1]
    return ax1, center

def __inverse_ellipse(axes, scalar=1):
    ax1 = 1/axes*axes[1]**2*scalar
    center = ax1[1]*N.sqrt(2)
    return ax1, center

def __reverse_ellipse(axes, scalar=1):
    """
    This method doesn't work as well
    """
    ax1 = axes.copy()[::-1]*scalar
    center = ax1[1]*N.sqrt(2)*scalar
    return ax1, center

def plot_conjugate_conics(ax, axes, width=None, plot_foci=False, plot_inverse_hyperbola=False):
    hyp, = ax.plot(*hyperbola(axes))

    if plot_inverse_hyperbola:
        # Plot inverse hyperbola
        hyp_inv, = ax.plot(*hyperbola(1/axes, opens_up=False))
        color = hyp_inv.get_color()
    else:
        color = None

    ax.plot(*ellipse(axes), zorder=-5, color=hyp.get_color(),alpha=0.5, linewidth=1)

    # Normal vector ellipse axes lengths
    # scales inversely with axes but can
    # be rescaled at will

    ax1,center = __inverse_ellipse(axes)

    ell, = ax.plot(*ellipse(ax1, center=[0,center]), color=color)

    if plot_foci:
        # Plot hyperbola focus
        hyp_c = N.sqrt(N.sum(axes**2))
        ax.plot([0,0],[hyp_c,-hyp_c], '.', color=hyp.get_color())

        # Plot ellipse foci
        c = ax1**2
        c.sort()
        c[0] *= -1
        ell_c = N.sqrt(N.sum(c))
        # Plot ellipse foci
        ax.plot([0,0],[center+ell_c,center-ell_c], '.', color=ell.get_color())

    # Plot tangents
    xvals = N.array([-500,500])
    yvals = axes[1]/axes[0]*xvals
    kw = dict(zorder=-1, color=hyp.get_color(), linewidth=0.5)
    ax.plot(xvals, yvals, ':', **kw)
    ax.plot(xvals, -yvals, ':', **kw)

    #yvals = axes[0]/axes[1]*xvals
    kw = dict(zorder=-1, color=ell.get_color(), linewidth=0.5)
    ax.plot(yvals, xvals, ':', **kw)
    ax.plot(yvals, -xvals, ':', **kw)

    _ = 4
    if axes[1] > 0.5*axes[0]:
        _ = 6

    if width is None:
        width = N.linalg.norm(axes)*_
    lim = N.array([-width,width])/2
    ax.set_xlim(lim)
    ax.set_ylim(lim*0.6)
    ax.set_aspect(1)

    return ax
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import unittest

from gilded_rose import Item, GildedRose

class GildedRoseTest(unittest.TestCase):
    def test_foo(self):
        items = [Item("foo", 0, 0)]
        gilded_rose = GildedRose(items)
        gilded_rose.update_quality()
        self.assertEquals("fixme", items[0].name)

if __name__ == '__main__':
    unittest.main()
<commit_msg>Update test_gilded_rose.py<commit_after># -*- coding: utf-8 -*-
import unittest

from gilded_rose import Item, GildedRose

class GildedRoseTest(unittest.TestCase):
    def test_foo(self):
        items = [Item("foo", 0, 0)]
        gilded_rose = GildedRose(items)
        gilded_rose.update_quality()
        self.assertEquals("fixme", items[0].name)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>