<commit_before>from firedrake import split, LinearVariationalProblem, \
    LinearVariationalSolver, TestFunctions, TrialFunctions, \
    TestFunction, TrialFunction, lhs, rhs, DirichletBC, FacetNormal, \
    div, dx, jump, avg, dS_v, dS_h, inner, MixedFunctionSpace, dot, grad, \
    Function, MixedVectorSpaceBasis, VectorSpaceBasis
from firedrake.solving_utils import flatten_parameters

from gusto.configuration import DEBUG
from gusto.forcing import exner, exner_rho, exner_theta
from abc import ABCMeta, abstractmethod, abstractproperty

__all__ = ["CompressibleSolver", "IncompressibleSolver", "ShallowWaterSolver"]

class TimesteppingSolver(object, metaclass=ABCMeta):
    """
    Base class for timestepping linear solvers for Gusto.

    This is a dummy base class.

    :arg state: :class:`.State` object.
    :arg solver_parameters (optional): solver parameters
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    """

    def __init__(self, state, solver_parameters=None,
                 overwrite_solver_parameters=False):

        self.state = state

        if solver_parameters is not None:
            if not overwrite_solver_parameters:
                p = flatten_parameters(self.solver_parameters)
                p.update(flatten_parameters(solver_parameters))
                solver_parameters = p
            self.solver_parameters = solver_parameters

        if state.output.log_level == DEBUG:
            self.solver_parameters["ksp_monitor_true_residual"] = True

        # setup the solver
        self._setup_solver()

    @abstractproperty
    def solver_parameters(self):
        """Solver parameters for this solver"""
        pass

    @abstractmethod
    def solve(self):
        pass

class CompressibleSolver(TimesteppingSolver):
    """
    Timestepping linear solver object for the compressible equations
    in theta-pi formulation with prognostic variables u,rho,theta.

    This solver follows the following strategy:
    (1) Analytically eliminate theta (introduces error near topography)
    (2) Solve resulting system for (u,rho) using a Schur preconditioner
    (3) Reconstruct theta

    :arg state: a :class:`.State` object containing everything else.
    :arg quadrature degree: tuple (q_h, q_v) where q_h is the required
    quadrature degree in the horizontal direction and q_v is that in
    the vertical direction
    :arg solver_parameters (optional): solver parameters
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    :arg moisture (optional): list of names of moisture fields.
    """

    solver_parameters = {
        'pc_type': 'fieldsplit',
        'pc_fieldsplit_type': 'schur',
        'ksp_type': 'gmres',
        'ksp_max_it': 100,
        'ksp_gmres_restart': 50,
        'pc_fieldsplit_schur_fact_type': 'FULL',
        'pc_fieldsplit_schur_precondition': 'selfp',
        'fieldsplit_0': {'ksp_type': 'preonly',
                         'pc_type': 'bjacobi',
                         'sub_pc_type': 'ilu'},
        'fieldsplit_1': {'ksp_type': 'preonly',
                         'pc_type': 'gamg',
                         'mg_levels': {'ksp_type': 'chebyshev',
                                       'ksp_chebyshev_esteig': True,
                                       'ksp_max_it': 1,
                                       'pc_type': 'bjacobi',
                                       'sub_pc_type': 'ilu'}}
    }

    def __init__(self, state, quadrature_degree=None, solver_parameters=None,
                 overwrite_solver_parameters=False, moisture=None):

        self.moisture = moisture

        if quadrature_degree is not None:
            self.quadrature_degree = quadrature_degree
        else:
            dgspace = state.spaces("DG")
            if any(deg > 2 for deg in dgspace.ufl_element().degree()):
                state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
            self.quadrature_degree = (5, 5)

        super().__init__(state, solver_parameters, overwrite_solver_parameters)

    def _setup_solver(self):
        state = self.state      # just cutting down line length a bit
        dt = state.timestepping.dt
        beta = dt*state.timestepping.alpha
        cp = state.parameters.cp
        mu = state.mu
        Vu = state.spaces("HDiv")
        Vtheta = state.spaces("HDiv_v")
        Vrho = state.spaces("DG")

        # Split up the rhs vector (symbolically)
        u_in, rho_in, theta_in = split(state.xrhs)

        # Build the reduced function space for u,rho
        M = MixedFunctionSpace((Vu, Vrho))
        w, phi = TestFunctions(M)
        u, rho = TrialFunctions(M)

        n = FacetNormal(state.mesh)

        # Get background fields
        thetabar = state.fields("thetabar")
        rhobar = state.fields("rhobar")
        pibar = exner(thetabar, rhobar, state)
        pibar_rho = exner_rho(thetabar, rhobar, state)
        pibar_theta = exner_theta(thetabar, rhobar, state)

        # Analytical (approximate) elimination of theta
        k = state.k             # Upward pointing unit vector
        theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in

        # Only include theta' (rather than pi') in the vertical
        # component of the gradient

        # the pi prime term (here, bars are for mean and no bars are
        # for linear perturbations)

        pi = pibar_theta*theta + pibar_rho*rho

        # vertical projection
        def V(u):
            return k*inner(u, k)

        # specify degree for some terms as estimated degree is too large
        dxp = dx(degree=(self.quadrature_degree))
        dS_vp = dS_v(degree=(self.quadrature_degree))

        # add effect of density of water upon theta
        if self.moisture is not None:
            water_t = Function(Vtheta).assign(0.0)
            for water in self.moisture:
                water_t += self.state.fields(water)
            theta = theta / (1 + water_t)
            thetabar = thetabar / (1 + water_t)

        eqn = (
            inner(w, (u - u_in))*dx
            - beta*cp*div(theta*V(w))*pibar*dxp
            # following does nothing but is preserved in the comments
            # to remind us why (because V(w) is purely vertical.
            # + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v
            - beta*cp*div(thetabar*w)*pi*dxp
            + beta*cp*jump(thetabar*w, n)*avg(pi)*dS_vp
            + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
            + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
        )

        if mu is not None:
            eqn += dt*mu*inner(w, k)*inner(u, k)*dx
        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u rho solver
        self.urho = Function(M)

        # Boundary conditions (assumes extruded mesh)
        bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
               DirichletBC(M.sub(0), 0.0, "top")]

        # Solver for u, rho
        urho_problem = LinearVariationalProblem(
            aeqn, Leqn, self.urho, bcs=bcs)

        self.urho_solver = LinearVariationalSolver(urho_problem,
                                                   solver_parameters=self.solver_parameters,
                                                   options_prefix='ImplicitSolver')

        # Reconstruction of theta
        theta = TrialFunction(Vtheta)
        gamma = TestFunction(Vtheta)

        u, rho = self.urho.split()
        self.theta = Function(Vtheta)

        theta_eqn = gamma*(theta - theta_in +
                           dot(k, u)*dot(k, grad(thetabar))*beta)*dx

        theta_problem = LinearVariationalProblem(lhs(theta_eqn),
                                                 rhs(theta_eqn),
                                                 self.theta)
        self.theta_solver = LinearVariationalSolver(theta_problem,
                                                    options_prefix='thetabacksubstitution')

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.urho_solver.solve()

        u1, rho1 = self.urho.split()
        u, rho, theta = self.state.dy.split()
        u.assign(u1)
        rho.assign(rho1)

        self.theta_solver.solve()
        theta.assign(self.theta)

class IncompressibleSolver(TimesteppingSolver):
    """Timestepping linear solver object for the incompressible
    Boussinesq equations with prognostic variables u, p, b.

    This solver follows the following strategy:
    (1) Analytically eliminate b (introduces error near topography)
    (2) Solve resulting system for (u,p) using a block Hdiv preconditioner
    (3) Reconstruct b

    This currently requires a (parallel) direct solver so is probably
    a bit memory-hungry, we'll improve this with a hybridised solver
    soon.

    :arg state: a :class:`.State` object containing everything else.
    :arg L: the width of the domain, used in the preconditioner.
    :arg solver_parameters: (optional) Solver parameters.
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    """

    solver_parameters = {
        'ksp_type': 'gmres',
        'pc_type': 'fieldsplit',
        'pc_fieldsplit_type': 'additive',
        'fieldsplit_0': {'ksp_type': 'preonly',
                         'pc_type': 'lu',
                         'pc_factor_mat_solver_package': 'mumps'},
        'fieldsplit_1': {'ksp_type': 'preonly',
                         'pc_type': 'lu',
                         'pc_factor_mat_solver_package': 'mumps'}
    }

    def __init__(self, state, L, solver_parameters=None,
                 overwrite_solver_parameters=False):

        self.L = L
        super().__init__(state, solver_parameters, overwrite_solver_parameters)

    def _setup_solver(self):
        state = self.state      # just cutting down line length a bit
        dt = state.timestepping.dt
        beta = dt*state.timestepping.alpha
        mu = state.mu
        Vu = state.spaces("HDiv")
        Vb = state.spaces("HDiv_v")
        Vp = state.spaces("DG")

        # Split up the rhs vector (symbolically)
        u_in, p_in, b_in = split(state.xrhs)

        # Build the reduced function space for u,p
        M = MixedFunctionSpace((Vu, Vp))
        w, phi = TestFunctions(M)
        u, p = TrialFunctions(M)

        # Get background fields
        bbar = state.fields("bbar")

        # Analytical (approximate) elimination of theta
        k = state.k             # Upward pointing unit vector
        b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in

        # vertical projection
        def V(u):
            return k*inner(u, k)

        eqn = (
            inner(w, (u - u_in))*dx
            - beta*div(w)*p*dx
            - beta*inner(w, k)*b*dx
            + phi*div(u)*dx
        )

        if mu is not None:
            eqn += dt*mu*inner(w, k)*inner(u, k)*dx
        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u p solver
        self.up = Function(M)

        # Boundary conditions (assumes extruded mesh)
        bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
               DirichletBC(M.sub(0), 0.0, "top")]

        # preconditioner equation
        L = self.L
        Ap = (
            inner(w, u) + L*L*div(w)*div(u) +
            phi*p/L/L
        )*dx

        # Solver for u, p
        up_problem = LinearVariationalProblem(
            aeqn, Leqn, self.up, bcs=bcs, aP=Ap)

        nullspace = MixedVectorSpaceBasis(M,
                                          [M.sub(0),
                                           VectorSpaceBasis(constant=True)])

        self.up_solver = LinearVariationalSolver(up_problem,
                                                 solver_parameters=self.solver_parameters,
                                                 nullspace=nullspace)

        # Reconstruction of b
        b = TrialFunction(Vb)
        gamma = TestFunction(Vb)

        u, p = self.up.split()
        self.b = Function(Vb)

        b_eqn = gamma*(b - b_in +
                       dot(k, u)*dot(k, grad(bbar))*beta)*dx

        b_problem = LinearVariationalProblem(lhs(b_eqn),
                                             rhs(b_eqn),
                                             self.b)
        self.b_solver = LinearVariationalSolver(b_problem)

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.up_solver.solve()

        u1, p1 = self.up.split()
        u, p, b = self.state.dy.split()
        u.assign(u1)
        p.assign(p1)

        self.b_solver.solve()
        b.assign(self.b)

class ShallowWaterSolver(TimesteppingSolver):

    solver_parameters = {
        'ksp_type': 'preonly',
        'mat_type': 'matfree',
        'pc_type': 'python',
        'pc_python_type': 'firedrake.HybridizationPC',
        'hybridization': {'ksp_type': 'cg',
                          'pc_type': 'gamg',
                          'ksp_rtol': 1e-8,
                          'mg_levels': {'ksp_type': 'chebyshev',
                                        'ksp_max_it': 2,
                                        'pc_type': 'bjacobi',
                                        'sub_pc_type': 'ilu'},
                          # Broken residual construction
                          'hdiv_residual': {'ksp_type': 'cg',
                                            'pc_type': 'bjacobi',
                                            'sub_pc_type': 'ilu',
                                            'ksp_rtol': 1e-8},
                          # Projection step
                          'hdiv_projection': {'ksp_type': 'cg',
                                              'ksp_rtol': 1e-8}}
    }

    def _setup_solver(self):
        state = self.state
        H = state.parameters.H
        g = state.parameters.g
        beta = state.timestepping.dt*state.timestepping.alpha

        # Split up the rhs vector (symbolically)
        u_in, D_in = split(state.xrhs)

        W = state.W
        w, phi = TestFunctions(W)
        u, D = TrialFunctions(W)

        eqn = (
            inner(w, u) - beta*g*div(w)*D
            - inner(w, u_in)
            + phi*D + beta*H*phi*div(u)
            - phi*D_in
        )*dx

        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u rho solver
        self.uD = Function(W)

        # Solver for u, D
        uD_problem = LinearVariationalProblem(
            aeqn, Leqn, self.state.dy)

        self.uD_solver = LinearVariationalSolver(uD_problem,
                                                 solver_parameters=self.solver_parameters,
                                                 options_prefix='SWimplicit')

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.uD_solver.solve()
<commit_msg>Prune dead solver options<commit_after>from firedrake import split, LinearVariationalProblem, \
    LinearVariationalSolver, TestFunctions, TrialFunctions, \
    TestFunction, TrialFunction, lhs, rhs, DirichletBC, FacetNormal, \
    div, dx, jump, avg, dS_v, dS_h, inner, MixedFunctionSpace, dot, grad, \
    Function, MixedVectorSpaceBasis, VectorSpaceBasis
from firedrake.solving_utils import flatten_parameters

from gusto.configuration import DEBUG
from gusto.forcing import exner, exner_rho, exner_theta
from abc import ABCMeta, abstractmethod, abstractproperty

__all__ = ["CompressibleSolver", "IncompressibleSolver", "ShallowWaterSolver"]

class TimesteppingSolver(object, metaclass=ABCMeta):
    """
    Base class for timestepping linear solvers for Gusto.

    This is a dummy base class.

    :arg state: :class:`.State` object.
    :arg solver_parameters (optional): solver parameters
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    """

    def __init__(self, state, solver_parameters=None,
                 overwrite_solver_parameters=False):

        self.state = state

        if solver_parameters is not None:
            if not overwrite_solver_parameters:
                p = flatten_parameters(self.solver_parameters)
                p.update(flatten_parameters(solver_parameters))
                solver_parameters = p
            self.solver_parameters = solver_parameters

        if state.output.log_level == DEBUG:
            self.solver_parameters["ksp_monitor_true_residual"] = True

        # setup the solver
        self._setup_solver()

    @abstractproperty
    def solver_parameters(self):
        """Solver parameters for this solver"""
        pass

    @abstractmethod
    def solve(self):
        pass

class CompressibleSolver(TimesteppingSolver):
    """
    Timestepping linear solver object for the compressible equations
    in theta-pi formulation with prognostic variables u,rho,theta.

    This solver follows the following strategy:
    (1) Analytically eliminate theta (introduces error near topography)
    (2) Solve resulting system for (u,rho) using a Schur preconditioner
    (3) Reconstruct theta

    :arg state: a :class:`.State` object containing everything else.
    :arg quadrature degree: tuple (q_h, q_v) where q_h is the required
    quadrature degree in the horizontal direction and q_v is that in
    the vertical direction
    :arg solver_parameters (optional): solver parameters
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    :arg moisture (optional): list of names of moisture fields.
    """

    solver_parameters = {
        'pc_type': 'fieldsplit',
        'pc_fieldsplit_type': 'schur',
        'ksp_type': 'gmres',
        'ksp_max_it': 100,
        'ksp_gmres_restart': 50,
        'pc_fieldsplit_schur_fact_type': 'FULL',
        'pc_fieldsplit_schur_precondition': 'selfp',
        'fieldsplit_0': {'ksp_type': 'preonly',
                         'pc_type': 'bjacobi',
                         'sub_pc_type': 'ilu'},
        'fieldsplit_1': {'ksp_type': 'preonly',
                         'pc_type': 'gamg',
                         'mg_levels': {'ksp_type': 'chebyshev',
                                       'ksp_chebyshev_esteig': True,
                                       'ksp_max_it': 1,
                                       'pc_type': 'bjacobi',
                                       'sub_pc_type': 'ilu'}}
    }

    def __init__(self, state, quadrature_degree=None, solver_parameters=None,
                 overwrite_solver_parameters=False, moisture=None):

        self.moisture = moisture

        if quadrature_degree is not None:
            self.quadrature_degree = quadrature_degree
        else:
            dgspace = state.spaces("DG")
            if any(deg > 2 for deg in dgspace.ufl_element().degree()):
                state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
            self.quadrature_degree = (5, 5)

        super().__init__(state, solver_parameters, overwrite_solver_parameters)

    def _setup_solver(self):
        state = self.state      # just cutting down line length a bit
        dt = state.timestepping.dt
        beta = dt*state.timestepping.alpha
        cp = state.parameters.cp
        mu = state.mu
        Vu = state.spaces("HDiv")
        Vtheta = state.spaces("HDiv_v")
        Vrho = state.spaces("DG")

        # Split up the rhs vector (symbolically)
        u_in, rho_in, theta_in = split(state.xrhs)

        # Build the reduced function space for u,rho
        M = MixedFunctionSpace((Vu, Vrho))
        w, phi = TestFunctions(M)
        u, rho = TrialFunctions(M)

        n = FacetNormal(state.mesh)

        # Get background fields
        thetabar = state.fields("thetabar")
        rhobar = state.fields("rhobar")
        pibar = exner(thetabar, rhobar, state)
        pibar_rho = exner_rho(thetabar, rhobar, state)
        pibar_theta = exner_theta(thetabar, rhobar, state)

        # Analytical (approximate) elimination of theta
        k = state.k             # Upward pointing unit vector
        theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in

        # Only include theta' (rather than pi') in the vertical
        # component of the gradient

        # the pi prime term (here, bars are for mean and no bars are
        # for linear perturbations)

        pi = pibar_theta*theta + pibar_rho*rho

        # vertical projection
        def V(u):
            return k*inner(u, k)

        # specify degree for some terms as estimated degree is too large
        dxp = dx(degree=(self.quadrature_degree))
        dS_vp = dS_v(degree=(self.quadrature_degree))

        # add effect of density of water upon theta
        if self.moisture is not None:
            water_t = Function(Vtheta).assign(0.0)
            for water in self.moisture:
                water_t += self.state.fields(water)
            theta = theta / (1 + water_t)
            thetabar = thetabar / (1 + water_t)

        eqn = (
            inner(w, (u - u_in))*dx
            - beta*cp*div(theta*V(w))*pibar*dxp
            # following does nothing but is preserved in the comments
            # to remind us why (because V(w) is purely vertical.
            # + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v
            - beta*cp*div(thetabar*w)*pi*dxp
            + beta*cp*jump(thetabar*w, n)*avg(pi)*dS_vp
            + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
            + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
        )

        if mu is not None:
            eqn += dt*mu*inner(w, k)*inner(u, k)*dx
        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u rho solver
        self.urho = Function(M)

        # Boundary conditions (assumes extruded mesh)
        bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
               DirichletBC(M.sub(0), 0.0, "top")]

        # Solver for u, rho
        urho_problem = LinearVariationalProblem(
            aeqn, Leqn, self.urho, bcs=bcs)

        self.urho_solver = LinearVariationalSolver(urho_problem,
                                                   solver_parameters=self.solver_parameters,
                                                   options_prefix='ImplicitSolver')

        # Reconstruction of theta
        theta = TrialFunction(Vtheta)
        gamma = TestFunction(Vtheta)

        u, rho = self.urho.split()
        self.theta = Function(Vtheta)

        theta_eqn = gamma*(theta - theta_in +
                           dot(k, u)*dot(k, grad(thetabar))*beta)*dx

        theta_problem = LinearVariationalProblem(lhs(theta_eqn),
                                                 rhs(theta_eqn),
                                                 self.theta)
        self.theta_solver = LinearVariationalSolver(theta_problem,
                                                    options_prefix='thetabacksubstitution')

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.urho_solver.solve()

        u1, rho1 = self.urho.split()
        u, rho, theta = self.state.dy.split()
        u.assign(u1)
        rho.assign(rho1)

        self.theta_solver.solve()
        theta.assign(self.theta)

class IncompressibleSolver(TimesteppingSolver):
    """Timestepping linear solver object for the incompressible
    Boussinesq equations with prognostic variables u, p, b.

    This solver follows the following strategy:
    (1) Analytically eliminate b (introduces error near topography)
    (2) Solve resulting system for (u,p) using a block Hdiv preconditioner
    (3) Reconstruct b

    This currently requires a (parallel) direct solver so is probably
    a bit memory-hungry, we'll improve this with a hybridised solver
    soon.

    :arg state: a :class:`.State` object containing everything else.
    :arg L: the width of the domain, used in the preconditioner.
    :arg solver_parameters: (optional) Solver parameters.
    :arg overwrite_solver_parameters: boolean, if True use only the
    solver_parameters that have been passed in, if False then update
    the default solver parameters with the solver_parameters passed in.
    """

    solver_parameters = {
        'ksp_type': 'gmres',
        'pc_type': 'fieldsplit',
        'pc_fieldsplit_type': 'additive',
        'fieldsplit_0': {'ksp_type': 'preonly',
                         'pc_type': 'lu',
                         'pc_factor_mat_solver_package': 'mumps'},
        'fieldsplit_1': {'ksp_type': 'preonly',
                         'pc_type': 'lu',
                         'pc_factor_mat_solver_package': 'mumps'}
    }

    def __init__(self, state, L, solver_parameters=None,
                 overwrite_solver_parameters=False):

        self.L = L
        super().__init__(state, solver_parameters, overwrite_solver_parameters)

    def _setup_solver(self):
        state = self.state      # just cutting down line length a bit
        dt = state.timestepping.dt
        beta = dt*state.timestepping.alpha
        mu = state.mu
        Vu = state.spaces("HDiv")
        Vb = state.spaces("HDiv_v")
        Vp = state.spaces("DG")

        # Split up the rhs vector (symbolically)
        u_in, p_in, b_in = split(state.xrhs)

        # Build the reduced function space for u,p
        M = MixedFunctionSpace((Vu, Vp))
        w, phi = TestFunctions(M)
        u, p = TrialFunctions(M)

        # Get background fields
        bbar = state.fields("bbar")

        # Analytical (approximate) elimination of theta
        k = state.k             # Upward pointing unit vector
        b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in

        # vertical projection
        def V(u):
            return k*inner(u, k)

        eqn = (
            inner(w, (u - u_in))*dx
            - beta*div(w)*p*dx
            - beta*inner(w, k)*b*dx
            + phi*div(u)*dx
        )

        if mu is not None:
            eqn += dt*mu*inner(w, k)*inner(u, k)*dx
        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u p solver
        self.up = Function(M)

        # Boundary conditions (assumes extruded mesh)
        bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
               DirichletBC(M.sub(0), 0.0, "top")]

        # preconditioner equation
        L = self.L
        Ap = (
            inner(w, u) + L*L*div(w)*div(u) +
            phi*p/L/L
        )*dx

        # Solver for u, p
        up_problem = LinearVariationalProblem(
            aeqn, Leqn, self.up, bcs=bcs, aP=Ap)

        nullspace = MixedVectorSpaceBasis(M,
                                          [M.sub(0),
                                           VectorSpaceBasis(constant=True)])

        self.up_solver = LinearVariationalSolver(up_problem,
                                                 solver_parameters=self.solver_parameters,
                                                 nullspace=nullspace)

        # Reconstruction of b
        b = TrialFunction(Vb)
        gamma = TestFunction(Vb)

        u, p = self.up.split()
        self.b = Function(Vb)

        b_eqn = gamma*(b - b_in +
                       dot(k, u)*dot(k, grad(bbar))*beta)*dx

        b_problem = LinearVariationalProblem(lhs(b_eqn),
                                             rhs(b_eqn),
                                             self.b)
        self.b_solver = LinearVariationalSolver(b_problem)

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.up_solver.solve()

        u1, p1 = self.up.split()
        u, p, b = self.state.dy.split()
        u.assign(u1)
        p.assign(p1)

        self.b_solver.solve()
        b.assign(self.b)

class ShallowWaterSolver(TimesteppingSolver):

    solver_parameters = {
        'ksp_type': 'preonly',
        'mat_type': 'matfree',
        'pc_type': 'python',
        'pc_python_type': 'firedrake.HybridizationPC',
        'hybridization': {'ksp_type': 'cg',
                          'pc_type': 'gamg',
                          'ksp_rtol': 1e-8,
                          'mg_levels': {'ksp_type': 'chebyshev',
                                        'ksp_max_it': 2,
                                        'pc_type': 'bjacobi',
                                        'sub_pc_type': 'ilu'}}
    }

    def _setup_solver(self):
        state = self.state
        H = state.parameters.H
        g = state.parameters.g
        beta = state.timestepping.dt*state.timestepping.alpha

        # Split up the rhs vector (symbolically)
        u_in, D_in = split(state.xrhs)

        W = state.W
        w, phi = TestFunctions(W)
        u, D = TrialFunctions(W)

        eqn = (
            inner(w, u) - beta*g*div(w)*D
            - inner(w, u_in)
            + phi*D + beta*H*phi*div(u)
            - phi*D_in
        )*dx

        aeqn = lhs(eqn)
        Leqn = rhs(eqn)

        # Place to put result of u rho solver
        self.uD = Function(W)

        # Solver for u, D
        uD_problem = LinearVariationalProblem(
            aeqn, Leqn, self.state.dy)

        self.uD_solver = LinearVariationalSolver(uD_problem,
                                                 solver_parameters=self.solver_parameters,
                                                 options_prefix='SWimplicit')

    def solve(self):
        """
        Apply the solver with rhs state.xrhs and result state.dy.
        """

        self.uD_solver.solve()
<|endoftext|>


<commit_before>"""
    Saucebrush Filters are callables that take a Saucebrush source and yield
    back filtered records.

    The Filter, YieldFilter, and FieldFilter abstract base types are provided
    for convenience.  Derived classes only need to implement process_record
    (or process_field for FieldFilter).
"""

from saucebrush import utils
import re
import time

######################
## Abstract Filters ##
######################

class Filter(object):
    """ ABC for filters that operate on records.

        All derived filters must provide a process_record(self, record) that
        takes a single record (python dictionary) and returns a result.
    """

    def process_record(self, record):
        """ Abstract method to be overridden.

            Called with a single record, should return modified record.
        """
        raise NotImplementedError('process_record not defined in ' +
                                  self.__class__.__name__)

    def reject_record(self, record, message):
        recipe = getattr(self, '_recipe')
        if recipe:
            recipe.reject_record(record, message)

    def attach(self, source, recipe=None):
        self._recipe = recipe
        for record in source:
            result = self.process_record(record)
            if result is not None:
                yield result

class YieldFilter(Filter):
    """ ABC for defining filters where process_record yields.

        If process_record cannot return exactly one result for every record
        it is passed, it should yield back as many records as needed and the
        filter must derive from YieldFilter.
    """

    def attach(self, source, recipe=None):
        self._recipe = recipe
        for record in source:
            for result in self.process_record(record):
                yield result

class FieldFilter(Filter):
    """ ABC for filters that do a single operation on individual fields.

        All derived filters must provide a process_field(self, item) that
        returns a modified item.  process_field is called on one or more keys
        passed into __init__.
    """

    def __init__(self, keys):
        super(FieldFilter, self).__init__()
        self._target_keys = utils.str_or_list(keys)

    def process_record(self, record):
        """ Calls process_field on all keys passed to __init__. """

        for key in self._target_keys:
            try:
                item = record[key]
                record[key] = self.process_field(item)
            except KeyError:
                # probably want to have a boolean to flag missing fields
                pass
        return record

    def process_field(self, item):
        """ Given a value, return the value that it should be replaced with. """

        raise NotImplementedError('process_field not defined in ' +
                                  self.__class__.__name__)

    def __unicode__(self):
        return '%s( %s )' % (self.__class__.__name__, str(self._target_keys))

class ConditionalFilter(YieldFilter):
    """ ABC for filters that only pass through records meeting a condition.

        All derived filters must provide a test_record(self, record) that
        returns True or False -- True indicating that the record should be
        passed through, and False preventing pass through.
    """
    def process_record(self, record):
        """ Yields all records for which self.test_record is true """

        if self.test_record(record):
            yield record

    def test_record(self, record):
        """ Given a record, return True iff it should be passed on """
        raise NotImplementedError('test_record not defined in ' +
                                  self.__class__.__name__)

#####################
## Generic Filters ##
#####################

class FieldModifier(FieldFilter):
    """ Filter that calls a given function on a given set of fields.

        FieldModifier(('spam','eggs'), abs) to call the abs method on the spam
        and eggs fields in each record filtered.
    """

    def __init__(self, keys, func):
        super(FieldModifier, self).__init__(keys)
        self._filter_func = func

    def process_field(self, item):
        return self._filter_func(item)

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__,
                                 str(self._target_keys), str(self._filter_func))

class FieldRemover(Filter):
    """ Filter that removes a given set of fields.

        FieldRemover(('spam', 'eggs')) removes the spam and eggs fields from
        every record filtered.
    """

    def __init__(self, keys):
        super(FieldRemover, self).__init__()
        self._target_keys = utils.str_or_list(keys)

    def process_record(self, record):
        for key in self._target_keys:
            record.pop(key, None)
        return record

    def __unicode__(self):
        return '%s( %s )' % (self.__class__.__name__, str(self._target_keys))

class FieldMerger(Filter):
    """ Filter that merges a given set of fields using a supplied merge_func.

        Takes a mapping (dictionary of new_column:(from_col1,from_col2))

        FieldMerger({"bacon": ("spam", "eggs")}, operator.add) creates a new
        column bacon that is the result of spam+eggs
    """

    def __init__(self, mapping, merge_func, keep_fields=False):
        super(FieldMerger, self).__init__()
        self._field_mapping = mapping
        self._merge_func = merge_func
        self._keep_fields = keep_fields

    def process_record(self, record):
        for to_col, from_cols in self._field_mapping.iteritems():
            if self._keep_fields:
                values = [record.get(col, None) for col in from_cols]
            else:
                values = [record.pop(col, None) for col in from_cols]
            record[to_col] = self._merge_func(*values)
        return record

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__,
                                 str(self._field_mapping),
                                 str(self._merge_func))

class FieldAdder(Filter):
    """ Filter that adds a new field.

        Takes a name for the new field and a value, field_value can be an
        iterable, a function, or a static value.

        from itertools import count
        FieldAdder('id', count)

        would yield a new column named id that uses the itertools count iterable
        to create sequentially numbered ids.
    """

    def __init__(self, field_name, field_value, replace=True):
        super(FieldAdder, self).__init__()
        self._field_name = field_name
        self._field_value = field_value
        if hasattr(self._field_value, '__iter__'):
            self._field_value = iter(self._field_value).next
        self._replace = replace

    def process_record(self, record):
        if self._field_name not in record or self._replace:
            if callable(self._field_value):
                record[self._field_name] = self._field_value()
            else:
                record[self._field_name] = self._field_value
        return record

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__, self._field_name,
                             str(self._field_value))

class FieldCopier(Filter):
    """ Filter that copies one field to another.

        Takes a dictionary mapping destination keys to source keys.

    """
    def __init__(self, copy_mapping):
        super(FieldCopier, self).__init__()
        self._copy_mapping = copy_mapping

    def process_record(self, record):
        # mapping is dest:source
        for dest, source in self._copy_mapping.iteritems():
            record[dest] = record[source]
        return record

class FieldRenamer(Filter):
    """ Filter that renames one field to another.

        Takes a dictionary mapping destination keys to source keys.
    """
    def __init__(self, rename_mapping):
        super(FieldRenamer, self).__init__()
        self._rename_mapping = rename_mapping

    def process_record(self, record):
        # mapping is dest:source
        for dest, source in self._rename_mapping.iteritems():
            record[dest] = record.pop(source)
        return record

class Splitter(Filter):
    """ Filter that splits nested data into different paths.

        Takes a dictionary of keys and a series of filters to run against the
        associated dictionaries.

        {'person': {'firstname': 'James', 'lastname': 'Turk'},
         'phones': [{'phone': '222-222-2222'}, {'phone': '335-333-3321'}]
        }
    """

    def __init__(self, split_mapping):
        super(Splitter, self).__init__()
        self._split_mapping = split_mapping

    def process_record(self, record):
        for key, filters in self._split_mapping.iteritems():

            # if the key doesn't exist -- move on to next key
            try:
                subrecord = record[key]
            except KeyError:
                continue

            # if a dict, use process_record directly
            if isinstance(subrecord, dict):
                for filter_ in filters:
                    subrecord = filter_.process_record(subrecord)

            # if a list or tuple, use attach
            elif isinstance(subrecord, (list, tuple)):
                for filter_ in filters:
                    subrecord = filter_.attach(subrecord, recipe=self._recipe)
                subrecord = [r for r in subrecord]  # unchain generators

            # place back from whence it came
            record[key] = subrecord
        return record

class Flattener(FieldFilter):
    """ Collapse a set of similar dictionaries into a list.

        Takes a dictionary of keys and flattens the key names:

        addresses = [{'addresses': [{'address': {'state':'NC', 'street':'146 shirley drive'}},
                            {'address': {'state':'NY', 'street':'3000 Winton Rd'}}]}]
        flattener = Flattener(['addresses'])

        would yield:

        {'addresses': [{'state': 'NC', 'street': '146 shirley drive'},
                       {'state': 'NY', 'street': '3000 Winton Rd'}]}
    """
    def __init__(self, keys):
        super(Flattener, self).__init__(keys)

    def process_field(self, item):
        result = []
        for d in item:
            rec = {}
            for values in d.values():
                rec.update(values)
            result.append(rec)
        return result

class DictFlattener(Filter):
    def __init__(self, keys, separator='_'):
        super(DictFlattener, self).__init__()
        self._keys = utils.str_or_list(keys)
        self._separator = separator

    def process_record(self, record):
        return utils.flatten(record, keys=self._keys, separator=self._separator)

class Unique(ConditionalFilter):
    """ Filter that ensures that all records passing through are unique.
    """

    def __init__(self):
        super(Unique, self).__init__()
        self._seen = set()

    def test_record(self, record):
        record_hash = hash(repr(record))
        if record_hash not in self._seen:
            self._seen.add(record_hash)
            return True
        else:
            return False

class UnicodeFilter(Filter):
    """ Convert all str elements in the record to Unicode.
    """

    def __init__(self, encoding='utf-8', errors='ignore'):
        super(UnicodeFilter, self).__init__()
        self._encoding = encoding
        self._errors = errors

    def process_record(self, record):
        for key, value in record.iteritems():
            if isinstance(value, str):
                record[key] = unicode(value, self._encoding, self._errors)
            elif isinstance(value, unicode):
                record[key] = value.decode(self._encoding, self._errors)
        return record

class StringFilter(Filter):

    def __init__(self, encoding='utf-8', errors='ignore'):
        super(StringFilter, self).__init__()
        self._encoding = encoding
        self._errors = errors

    def process_record(self, record):
        for key, value in record.iteritems():
            if isinstance(value, unicode):
                record[key] = value.encode(self._encoding, self._errors)
        return record

###########################
## Commonly Used Filters ##
###########################

class PhoneNumberCleaner(FieldFilter):
    """ Filter that cleans phone numbers to match a given format.

        Takes a list of target keys and an optional phone # format that has
        10 %s placeholders.

        PhoneNumberCleaner( ('phone','fax'), number_format='%s%s%s-%s%s%s-%s%s%s%s')
        would format the phone & fax columns to 555-123-4567 format.
    """
    def __init__(self, keys, number_format='%s%s%s.%s%s%s.%s%s%s%s'):
        super(PhoneNumberCleaner, self).__init__(keys)
        self._number_format = number_format
        self._num_re = re.compile('\d')

    def process_field(self, item):
        nums = self._num_re.findall(item)
        if len(nums) == 10:
            item = self._number_format % tuple(nums)
        return item

class DateCleaner(FieldFilter):
    """ Filter that cleans dates to match a given format.

        Takes a list of target keys and to and from formats in strftime format.
    """
    def __init__(self, keys, from_format, to_format):
        super(DateCleaner, self).__init__(keys)
        self._from_format = from_format
        self._to_format = to_format

    def process_field(self, item):
        return time.strftime(self._to_format,
                             time.strptime(item, self._from_format))

class NameCleaner(Filter):
    """ Filter that splits names into a first, last, and middle name field.

        Takes a list of target keys.

        NameCleaner( ('name', ), nomatch_name='raw_name')
        would attempt to split 'name' into firstname, middlename, lastname,
        and suffix columns, and if it did not fit would place it in raw_name
    """

    # first middle? last suffix?
    FIRST_LAST = re.compile('''^\s*(?:(?P<firstname>\w+)(?:\.?)
                                \s+(?:(?P<middlename>\w+)\.?\s+)?
                                (?P<lastname>[A-Za-z'-]+))
                                (?:\s+(?P<suffix>JR\.?|II|III|IV))?
                                \s*$''', re.VERBOSE | re.IGNORECASE)

    # last, first middle? suffix?
    LAST_FIRST = re.compile('''^\s*(?:(?P<lastname>[A-Za-z'-]+),
                                \s+(?P<firstname>\w+)(?:\.?)
                                (?:\s+(?P<middlename>\w+)\.?)?)
                                (?:\s+(?P<suffix>JR\.?|II|III|IV))?
                                \s*$''', re.VERBOSE | re.IGNORECASE)

    def __init__(self, keys, prefix='', formats=None, nomatch_name=None):
        super(NameCleaner, self).__init__()
        self._keys = utils.str_or_list(keys)
        self._name_prefix = prefix
        self._nomatch_name = nomatch_name
        if formats:
            self._name_formats = formats
        else:
            self._name_formats = [self.FIRST_LAST, self.LAST_FIRST]

    def process_record(self, record):
        # run for each key (not using a FieldFilter due to multi-field output)
        for key in self._keys:
            name = record[key]

            # check if key matches any formats
            for format in self._name_formats:
                match = format.match(name)

                # if there is a match, remove original name and add pieces
                if match:
                    record.pop(key)
                    for k,v in match.groupdict().iteritems():
                        record[self._name_prefix + k] = v
                    break

            # if there is no match, move name into nomatch_name
            else:
                if self._nomatch_name:
                    record.pop(key)
                    record[self._nomatch_name] = name

        return record
<commit_msg>added UniqueID filter<commit_after>"""
    Saucebrush Filters are callables that take a Saucebrush source and yield
    back filtered records.

    The Filter, YieldFilter, and FieldFilter abstract base types are provided
    for convenience.  Derived classes only need to implement process_record
    (or process_field for FieldFilter).
"""

from saucebrush import utils
import re
import time

######################
## Abstract Filters ##
######################

class Filter(object):
    """ ABC for filters that operate on records.

        All derived filters must provide a process_record(self, record) that
        takes a single record (python dictionary) and returns a result.
    """

    def process_record(self, record):
        """ Abstract method to be overridden.

            Called with a single record, should return modified record.
        """
        raise NotImplementedError('process_record not defined in ' +
                                  self.__class__.__name__)

    def reject_record(self, record, message):
        recipe = getattr(self, '_recipe')
        if recipe:
            recipe.reject_record(record, message)

    def attach(self, source, recipe=None):
        self._recipe = recipe
        for record in source:
            result = self.process_record(record)
            if result is not None:
                yield result

class YieldFilter(Filter):
    """ ABC for defining filters where process_record yields.

        If process_record cannot return exactly one result for every record
        it is passed, it should yield back as many records as needed and the
        filter must derive from YieldFilter.
    """

    def attach(self, source, recipe=None):
        self._recipe = recipe
        for record in source:
            for result in self.process_record(record):
                yield result

class FieldFilter(Filter):
    """ ABC for filters that do a single operation on individual fields.

        All derived filters must provide a process_field(self, item) that
        returns a modified item.  process_field is called on one or more keys
        passed into __init__.
    """

    def __init__(self, keys):
        super(FieldFilter, self).__init__()
        self._target_keys = utils.str_or_list(keys)

    def process_record(self, record):
        """ Calls process_field on all keys passed to __init__. """

        for key in self._target_keys:
            try:
                item = record[key]
                record[key] = self.process_field(item)
            except KeyError:
                # probably want to have a boolean to flag missing fields
                pass
        return record

    def process_field(self, item):
        """ Given a value, return the value that it should be replaced with. """

        raise NotImplementedError('process_field not defined in ' +
                                  self.__class__.__name__)

    def __unicode__(self):
        return '%s( %s )' % (self.__class__.__name__, str(self._target_keys))

class ConditionalFilter(YieldFilter):
    """ ABC for filters that only pass through records meeting a condition.

        All derived filters must provide a test_record(self, record) that
        returns True or False -- True indicating that the record should be
        passed through, and False preventing pass through.
    """
    def process_record(self, record):
        """ Yields all records for which self.test_record is true """

        if self.test_record(record):
            yield record

    def test_record(self, record):
        """ Given a record, return True iff it should be passed on """
        raise NotImplementedError('test_record not defined in ' +
                                  self.__class__.__name__)

#####################
## Generic Filters ##
#####################

class FieldModifier(FieldFilter):
    """ Filter that calls a given function on a given set of fields.

        FieldModifier(('spam','eggs'), abs) to call the abs method on the spam
        and eggs fields in each record filtered.
    """

    def __init__(self, keys, func):
        super(FieldModifier, self).__init__(keys)
        self._filter_func = func

    def process_field(self, item):
        return self._filter_func(item)

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__,
                                 str(self._target_keys), str(self._filter_func))

class FieldRemover(Filter):
    """ Filter that removes a given set of fields.

        FieldRemover(('spam', 'eggs')) removes the spam and eggs fields from
        every record filtered.
    """

    def __init__(self, keys):
        super(FieldRemover, self).__init__()
        self._target_keys = utils.str_or_list(keys)

    def process_record(self, record):
        for key in self._target_keys:
            record.pop(key, None)
        return record

    def __unicode__(self):
        return '%s( %s )' % (self.__class__.__name__, str(self._target_keys))

class FieldMerger(Filter):
    """ Filter that merges a given set of fields using a supplied merge_func.

        Takes a mapping (dictionary of new_column:(from_col1,from_col2))

        FieldMerger({"bacon": ("spam", "eggs")}, operator.add) creates a new
        column bacon that is the result of spam+eggs
    """

    def __init__(self, mapping, merge_func, keep_fields=False):
        super(FieldMerger, self).__init__()
        self._field_mapping = mapping
        self._merge_func = merge_func
        self._keep_fields = keep_fields

    def process_record(self, record):
        for to_col, from_cols in self._field_mapping.iteritems():
            if self._keep_fields:
                values = [record.get(col, None) for col in from_cols]
            else:
                values = [record.pop(col, None) for col in from_cols]
            record[to_col] = self._merge_func(*values)
        return record

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__,
                                 str(self._field_mapping),
                                 str(self._merge_func))

class FieldAdder(Filter):
    """ Filter that adds a new field.

        Takes a name for the new field and a value, field_value can be an
        iterable, a function, or a static value.

        from itertools import count
        FieldAdder('id', count)

        would yield a new column named id that uses the itertools count iterable
        to create sequentially numbered ids.
    """

    def __init__(self, field_name, field_value, replace=True):
        super(FieldAdder, self).__init__()
        self._field_name = field_name
        self._field_value = field_value
        if hasattr(self._field_value, '__iter__'):
            self._field_value = iter(self._field_value).next
        self._replace = replace

    def process_record(self, record):
        if self._field_name not in record or self._replace:
            if callable(self._field_value):
                record[self._field_name] = self._field_value()
            else:
                record[self._field_name] = self._field_value
        return record

    def __unicode__(self):
        return '%s( %s, %s )' % (self.__class__.__name__, self._field_name,
                             str(self._field_value))

class FieldCopier(Filter):
    """ Filter that copies one field to another.

        Takes a dictionary mapping destination keys to source keys.

    """
    def __init__(self, copy_mapping):
        super(FieldCopier, self).__init__()
        self._copy_mapping = copy_mapping

    def process_record(self, record):
        # mapping is dest:source
        for dest, source in self._copy_mapping.iteritems():
            record[dest] = record[source]
        return record

class FieldRenamer(Filter):
    """ Filter that renames one field to another.

        Takes a dictionary mapping destination keys to source keys.
    """
    def __init__(self, rename_mapping):
        super(FieldRenamer, self).__init__()
        self._rename_mapping = rename_mapping

    def process_record(self, record):
        # mapping is dest:source
        for dest, source in self._rename_mapping.iteritems():
            record[dest] = record.pop(source)
        return record

class Splitter(Filter):
    """ Filter that splits nested data into different paths.

        Takes a dictionary of keys and a series of filters to run against the
        associated dictionaries.

        {'person': {'firstname': 'James', 'lastname': 'Turk'},
         'phones': [{'phone': '222-222-2222'}, {'phone': '335-333-3321'}]
        }
    """

    def __init__(self, split_mapping):
        super(Splitter, self).__init__()
        self._split_mapping = split_mapping

    def process_record(self, record):
        for key, filters in self._split_mapping.iteritems():

            # if the key doesn't exist -- move on to next key
            try:
                subrecord = record[key]
            except KeyError:
                continue

            # if a dict, use process_record directly
            if isinstance(subrecord, dict):
                for filter_ in filters:
                    subrecord = filter_.process_record(subrecord)

            # if a list or tuple, use attach
            elif isinstance(subrecord, (list, tuple)):
                for filter_ in filters:
                    subrecord = filter_.attach(subrecord, recipe=self._recipe)
                subrecord = [r for r in subrecord]  # unchain generators

            # place back from whence it came
            record[key] = subrecord
        return record

class Flattener(FieldFilter):
    """ Collapse a set of similar dictionaries into a list.

        Takes a dictionary of keys and flattens the key names:

        addresses = [{'addresses': [{'address': {'state':'NC', 'street':'146 shirley drive'}},
                            {'address': {'state':'NY', 'street':'3000 Winton Rd'}}]}]
        flattener = Flattener(['addresses'])

        would yield:

        {'addresses': [{'state': 'NC', 'street': '146 shirley drive'},
                       {'state': 'NY', 'street': '3000 Winton Rd'}]}
    """
    def __init__(self, keys):
        super(Flattener, self).__init__(keys)

    def process_field(self, item):
        result = []
        for d in item:
            rec = {}
            for values in d.values():
                rec.update(values)
            result.append(rec)
        return result

class DictFlattener(Filter):
    def __init__(self, keys, separator='_'):
        super(DictFlattener, self).__init__()
        self._keys = utils.str_or_list(keys)
        self._separator = separator

    def process_record(self, record):
        return utils.flatten(record, keys=self._keys, separator=self._separator)

class Unique(ConditionalFilter):
    """ Filter that ensures that all records passing through are unique.
    """

    def __init__(self):
        super(Unique, self).__init__()
        self._seen = set()

    def test_record(self, record):
        record_hash = hash(repr(record))
        if record_hash not in self._seen:
            self._seen.add(record_hash)
            return True
        else:
            return False

class UniqueID(ConditionalFilter):
    """ Filter that ensures that all records through have a unique ID.

        Takes the name of an ID field, or multiple field names in the case
        of a composite ID.
    """

    def __init__(self, field='id', *args):
        super(UniqueID, self).__init__()
        self._seen = set()
        self._id_fields = [field]
        self._id_fields.extend(args)

    def test_record(self, record):
        id_hash = hash(repr([record[key] for key in self._id_fields]))
        if id_hash not in self._seen:
            self._seen.add(id_hash)
            return True
        else:
            return False

class UnicodeFilter(Filter):
    """ Convert all str elements in the record to Unicode.
    """

    def __init__(self, encoding='utf-8', errors='ignore'):
        super(UnicodeFilter, self).__init__()
        self._encoding = encoding
        self._errors = errors

    def process_record(self, record):
        for key, value in record.iteritems():
            if isinstance(value, str):
                record[key] = unicode(value, self._encoding, self._errors)
            elif isinstance(value, unicode):
                record[key] = value.decode(self._encoding, self._errors)
        return record

class StringFilter(Filter):

    def __init__(self, encoding='utf-8', errors='ignore'):
        super(StringFilter, self).__init__()
        self._encoding = encoding
        self._errors = errors

    def process_record(self, record):
        for key, value in record.iteritems():
            if isinstance(value, unicode):
                record[key] = value.encode(self._encoding, self._errors)
        return record

###########################
## Commonly Used Filters ##
###########################

class PhoneNumberCleaner(FieldFilter):
    """ Filter that cleans phone numbers to match a given format.

        Takes a list of target keys and an optional phone # format that has
        10 %s placeholders.

        PhoneNumberCleaner( ('phone','fax'), number_format='%s%s%s-%s%s%s-%s%s%s%s')
        would format the phone & fax columns to 555-123-4567 format.
    """
    def __init__(self, keys, number_format='%s%s%s.%s%s%s.%s%s%s%s'):
        super(PhoneNumberCleaner, self).__init__(keys)
        self._number_format = number_format
        self._num_re = re.compile('\d')

    def process_field(self, item):
        nums = self._num_re.findall(item)
        if len(nums) == 10:
            item = self._number_format % tuple(nums)
        return item

class DateCleaner(FieldFilter):
    """ Filter that cleans dates to match a given format.

        Takes a list of target keys and to and from formats in strftime format.
    """
    def __init__(self, keys, from_format, to_format):
        super(DateCleaner, self).__init__(keys)
        self._from_format = from_format
        self._to_format = to_format

    def process_field(self, item):
        return time.strftime(self._to_format,
                             time.strptime(item, self._from_format))

class NameCleaner(Filter):
    """ Filter that splits names into a first, last, and middle name field.

        Takes a list of target keys.

        NameCleaner( ('name', ), nomatch_name='raw_name')
        would attempt to split 'name' into firstname, middlename, lastname,
        and suffix columns, and if it did not fit would place it in raw_name
    """

    # first middle? last suffix?
    FIRST_LAST = re.compile('''^\s*(?:(?P<firstname>\w+)(?:\.?)
                                \s+(?:(?P<middlename>\w+)\.?\s+)?
                                (?P<lastname>[A-Za-z'-]+))
                                (?:\s+(?P<suffix>JR\.?|II|III|IV))?
                                \s*$''', re.VERBOSE | re.IGNORECASE)

    # last, first middle? suffix?
    LAST_FIRST = re.compile('''^\s*(?:(?P<lastname>[A-Za-z'-]+),
                                \s+(?P<firstname>\w+)(?:\.?)
                                (?:\s+(?P<middlename>\w+)\.?)?)
                                (?:\s+(?P<suffix>JR\.?|II|III|IV))?
                                \s*$''', re.VERBOSE | re.IGNORECASE)

    def __init__(self, keys, prefix='', formats=None, nomatch_name=None):
        super(NameCleaner, self).__init__()
        self._keys = utils.str_or_list(keys)
        self._name_prefix = prefix
        self._nomatch_name = nomatch_name
        if formats:
            self._name_formats = formats
        else:
            self._name_formats = [self.FIRST_LAST, self.LAST_FIRST]

    def process_record(self, record):
        # run for each key (not using a FieldFilter due to multi-field output)
        for key in self._keys:
            name = record[key]

            # check if key matches any formats
            for format in self._name_formats:
                match = format.match(name)

                # if there is a match, remove original name and add pieces
                if match:
                    record.pop(key)
                    for k,v in match.groupdict().iteritems():
                        record[self._name_prefix + k] = v
                    break

            # if there is no match, move name into nomatch_name
            else:
                if self._nomatch_name:
                    record.pop(key)
                    record[self._nomatch_name] = name

        return record
<|endoftext|>


<commit_before><commit_msg>Bugfix to Lyons.<commit_after><|endoftext|>


<commit_before>#! /usr/bin/env python
"""Unit tests for smartcard.ATR

This test case can be executed individually, or with all other test cases
thru testsuite_framework.py.

__author__ = "http://www.gemalto.com"

Copyright 2001-2011 gemalto
Author: Jean-Daniel Aussel, mailto:jean-daniel.aussel@gemalto.com

This file is part of pyscard.

pyscard is free software; you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation; either version 2.1 of the License, or
(at your option) any later version.

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

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

import unittest
import string

# gemalto jython
from smartcard.ATR import ATR

class testcase_CAtr(unittest.TestCase):
    """Test APDU class and utilities"""

    def testcase_ATR1(self):
        """Usimera Classic 2."""
        a = ATR([0x3B, 0x9E, 0x95, 0x80, 0x1F, 0xC3, 0x80, 0x31, 0xA0, 0x73,
            0xBE, 0x21, 0x13, 0x67, 0x29, 0x02, 0x01, 0x01, 0x81, 0xCD, 0xB9])
        historicalbytes = [0x80, 0x31, 0xA0, 0x73, 0xBE, 0x21, 0x13, 0x67,
            0x29, 0x02, 0x01, 0x01, 0x81, 0xCD]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), 0xB9)
        self.assert_(a.checksumOK)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(a.isT15Supported())

    def testcase_ATR2(self):
        """Palmera Protect V2."""
        a = ATR([0x3B, 0x65, 0x00, 0x00, 0x9C, 0x02, 0x02, 0x01, 0x02])
        historicalbytes = [0x9C, 0x02, 0x02, 0x01, 0x02]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR3(self):
        """Simera 3.13."""
        a = ATR([0x3B, 0x16, 0x18, 0x20, 0x02, 0x01, 0x00, 0x80, 0x0D])
        historicalbytes = [0x20, 0x02, 0x01, 0x00, 0x80, 0x0D]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR4(self):
        """SIMRock'n Tree"""
        a = ATR([0x3B, 0x77, 0x94, 0x00, 0x00, 0x82, 0x30, 0x00, 0x13,
            0x6C, 0x9F, 0x22])
        historicalbytes = [0x82, 0x30, 0x00, 0x13, 0x6C, 0x9F, 0x22]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR5(self):
        """Demo Vitale online IGEA340"""
        a = ATR([0x3F, 0x65, 0x25, 0x00, 0x52, 0x09, 0x6A, 0x90, 0x00])
        historicalbytes = [0x52, 0x09, 0x6A, 0x90, 0x00]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR6(self):
        """Simagine 2002"""
        a = ATR([0x3B, 0x16, 0x94, 0x20, 0x02, 0x01, 0x00, 0x00, 0x0D])
        historicalbytes = [0x20, 0x02, 0x01, 0x00, 0x00, 0x0D]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)

    def testcase_ATR7(self):
        """Protect V3 T=1"""
        a = ATR([0x3B, 0xE5, 0x00, 0x00, 0x81, 0x21, 0x45, 0x9C, 0x10,
            0x01, 0x00, 0x80, 0x0D])
        historicalbytes = [0x9C, 0x10, 0x01, 0x00, 0x80]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), 0x0D)
        self.assert_(not a.isT0Supported())
        self.assert_(a.isT1Supported())
        self.assert_(not a.isT15Supported())
        self.assert_(a.checksumOK)
        self.assert_(a.getTB1() == 0x00)
        self.assert_(a.getTC1() == 0x00)
        self.assert_(a.getTD1() == 0x81)
        self.assert_(a.TD[2 - 1] == 0x21) # TD2
        self.assert_(a.TB[3 - 1] == 0x45) # TB3

def suite():
    suite1 = unittest.makeSuite(testcase_CAtr)
    return unittest.TestSuite((suite1))

if __name__ == '__main__':
    unittest.main()
<commit_msg>Fixed E261 pep8 error at least two spaces before inline commen<commit_after>#! /usr/bin/env python
"""Unit tests for smartcard.ATR

This test case can be executed individually, or with all other test cases
thru testsuite_framework.py.

__author__ = "http://www.gemalto.com"

Copyright 2001-2011 gemalto
Author: Jean-Daniel Aussel, mailto:jean-daniel.aussel@gemalto.com

This file is part of pyscard.

pyscard is free software; you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation; either version 2.1 of the License, or
(at your option) any later version.

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

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

import unittest
import string

# gemalto jython
from smartcard.ATR import ATR

class testcase_CAtr(unittest.TestCase):
    """Test APDU class and utilities"""

    def testcase_ATR1(self):
        """Usimera Classic 2."""
        a = ATR([0x3B, 0x9E, 0x95, 0x80, 0x1F, 0xC3, 0x80, 0x31, 0xA0, 0x73,
            0xBE, 0x21, 0x13, 0x67, 0x29, 0x02, 0x01, 0x01, 0x81, 0xCD, 0xB9])
        historicalbytes = [0x80, 0x31, 0xA0, 0x73, 0xBE, 0x21, 0x13, 0x67,
            0x29, 0x02, 0x01, 0x01, 0x81, 0xCD]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), 0xB9)
        self.assert_(a.checksumOK)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(a.isT15Supported())

    def testcase_ATR2(self):
        """Palmera Protect V2."""
        a = ATR([0x3B, 0x65, 0x00, 0x00, 0x9C, 0x02, 0x02, 0x01, 0x02])
        historicalbytes = [0x9C, 0x02, 0x02, 0x01, 0x02]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR3(self):
        """Simera 3.13."""
        a = ATR([0x3B, 0x16, 0x18, 0x20, 0x02, 0x01, 0x00, 0x80, 0x0D])
        historicalbytes = [0x20, 0x02, 0x01, 0x00, 0x80, 0x0D]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR4(self):
        """SIMRock'n Tree"""
        a = ATR([0x3B, 0x77, 0x94, 0x00, 0x00, 0x82, 0x30, 0x00, 0x13,
            0x6C, 0x9F, 0x22])
        historicalbytes = [0x82, 0x30, 0x00, 0x13, 0x6C, 0x9F, 0x22]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR5(self):
        """Demo Vitale online IGEA340"""
        a = ATR([0x3F, 0x65, 0x25, 0x00, 0x52, 0x09, 0x6A, 0x90, 0x00])
        historicalbytes = [0x52, 0x09, 0x6A, 0x90, 0x00]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)
        self.assert_(a.isT0Supported())
        self.assert_(not a.isT1Supported())
        self.assert_(not a.isT15Supported())

    def testcase_ATR6(self):
        """Simagine 2002"""
        a = ATR([0x3B, 0x16, 0x94, 0x20, 0x02, 0x01, 0x00, 0x00, 0x0D])
        historicalbytes = [0x20, 0x02, 0x01, 0x00, 0x00, 0x0D]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), None)

    def testcase_ATR7(self):
        """Protect V3 T=1"""
        a = ATR([0x3B, 0xE5, 0x00, 0x00, 0x81, 0x21, 0x45, 0x9C, 0x10,
            0x01, 0x00, 0x80, 0x0D])
        historicalbytes = [0x9C, 0x10, 0x01, 0x00, 0x80]
        self.assertEquals(a.getHistoricalBytes(), historicalbytes)
        self.assertEquals(a.getChecksum(), 0x0D)
        self.assert_(not a.isT0Supported())
        self.assert_(a.isT1Supported())
        self.assert_(not a.isT15Supported())
        self.assert_(a.checksumOK)
        self.assert_(a.getTB1() == 0x00)
        self.assert_(a.getTC1() == 0x00)
        self.assert_(a.getTD1() == 0x81)
        self.assert_(a.TD[2 - 1] == 0x21)  # TD2
        self.assert_(a.TB[3 - 1] == 0x45)  # TB3

def suite():
    suite1 = unittest.makeSuite(testcase_CAtr)
    return unittest.TestSuite((suite1))

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


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
    Class to interface with MSFT OneDrive api
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
from onedrivesdk import get_default_client, Folder, Item
try:
    import cPickle as pickle
except ImportError:
    import pickle

from sync_app.get_auth_code_server import get_auth_code
from sync_app.util import HOMEDIR

class OneDriveInstance(object):
    """ class to make use of google python api """

    def __init__(self, number_to_process=-1):
        """ init function """
        self.redirect_uri = ''
        self.client_id = ''
        self.client_secret = ''
        self.list_of_keys = {}
        self.list_of_mimetypes = {}
        self.items_processed = 0
        self.list_of_folders = {}
        self.list_of_items = {}
        self.number_to_process = number_to_process
        self.read_credentials()
        self.client = self.get_auth()

    def read_credentials(self,
                         credential_file=HOMEDIR+'/.onedrive/credentials'):
        """ read credentials from file """
        with open(credential_file, 'r') as credfile:
            for line in credfile:
                key_, val_ = line.split()[:2]
                for key in ('redirect_uri', 'client_id', 'client_secret'):
                    if key.lower() == key_.strip().lower():
                        setattr(self, key, val_)

    def get_auth(self):
        """ do authorization """
        self.client = get_default_client(client_id=self.client_id,
                                         scopes=['wl.signin',
                                                 'wl.offline_access',
                                                 'onedrive.readwrite'])
        if os.path.exists('.onedrive_session.pkl'):
            with open('.onedrive_session.pkl', 'rb') as pfile:
                self.client.auth_provider._session = pickle.load(pfile)
        else:
            auth_url = \
                self.client.auth_provider.get_auth_url(self.redirect_uri)
            code = get_auth_code(auth_url, self.redirect_uri)
            self.client.auth_provider.authenticate(code, self.redirect_uri,
                                                   self.client_secret)
            with open('.onedrive_session.pkl', 'wb') as pfile:
                pickle.dump(self.client.auth_provider._session, pfile)

        return self.client

    def list_files(self, callback_fn):
        """ list non-directory files """
        def walk_nodes(parentid='root'):
            parent_node = self.client.item(id=parentid)
            for node in parent_node.children.get():
                item = node.to_dict()
                item['parentid'] = parentid
                if 'folder' in item and item['folder']['childCount'] > 0:
                    walk_nodes(parentid=item['id'])
                else:
                    callback_fn(item)

        walk_nodes(parentid='root')

    def get_folders(self, callback_fn):
        """ get folders """
        def walk_nodes(parentid='root'):
            parent_node = self.client.item(id=parentid)
            for node in parent_node.children.get():
                item = node.to_dict()
                item['parentid'] = parentid
                if 'folder' in item:
                    callback_fn(item)
                    if item['folder']['childCount'] > 0:
                        walk_nodes(parentid=item['id'])
        walk_nodes(parentid='root')

    def download(self, did, exportfile, sha1sum=None):
        """ download using dlink url """
        dirname = os.path.dirname(exportfile)
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        self.client.item(id=did).download(exportfile + '.new')
        if sha1sum:
            from sync_app.util import get_sha1
            sha = get_sha1(exportfile + '.new')
            if sha != sha1sum:
                raise TypeError('%s %s' % (sha, sha1sum))
        os.rename('%s.new' % exportfile, exportfile)
        return True

    def upload(self, fname, parent_id='root'):
        """ upload fname and assign parent_id if provided """
        bname = os.path.basename(fname)
        node = self.client.item(drive="me", id=parent_id)
        return node.children[bname].upload(fname).to_dict()

    def create_directory(self, dname, parent_id='root'):
        """ create directory, assign parent_id if supplied """
        if not parent_id:
            raise ValueError('need to specify parent_id')
        newfolder = Folder()
        newitem = Item()
        newitem.name = dname
        newitem.folder = newfolder

        result = self.client.item(id=parent_id).children.add(newitem).to_dict()
        result['parentid'] = parent_id
        return result

    def delete_file(self, fileid):
        """ delete file by fileid """
        return self.client.item(id=fileid).delete()
<commit_msg>t_get<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""
    Class to interface with MSFT OneDrive api
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
import time
from onedrivesdk import get_default_client, Folder, Item
try:
    import cPickle as pickle
except ImportError:
    import pickle

from sync_app.get_auth_code_server import get_auth_code
from sync_app.util import HOMEDIR

def t_get(request):
    timeout = 1
    while True:
        try:
            return request.get()
        except Exception as exc:
            print('timeout %s, %s' % (timeout, exc))
            time.sleep(timeout)
            timeout *= 2
            if timeout >= 64:
                raise

class OneDriveInstance(object):
    """ class to make use of google python api """

    def __init__(self, number_to_process=-1):
        """ init function """
        self.redirect_uri = ''
        self.client_id = ''
        self.client_secret = ''
        self.list_of_keys = {}
        self.list_of_mimetypes = {}
        self.items_processed = 0
        self.list_of_folders = {}
        self.list_of_items = {}
        self.number_to_process = number_to_process
        self.read_credentials()
        self.client = self.get_auth()

    def read_credentials(self,
                         credential_file=HOMEDIR+'/.onedrive/credentials'):
        """ read credentials from file """
        with open(credential_file, 'r') as credfile:
            for line in credfile:
                key_, val_ = line.split()[:2]
                for key in ('redirect_uri', 'client_id', 'client_secret'):
                    if key.lower() == key_.strip().lower():
                        setattr(self, key, val_)

    def get_auth(self):
        """ do authorization """
        self.client = get_default_client(client_id=self.client_id,
                                         scopes=['wl.signin',
                                                 'wl.offline_access',
                                                 'onedrive.readwrite'])
        if os.path.exists('.onedrive_session.pkl'):
            with open('.onedrive_session.pkl', 'rb') as pfile:
                self.client.auth_provider._session = pickle.load(pfile)
        else:
            auth_url = \
                self.client.auth_provider.get_auth_url(self.redirect_uri)
            code = get_auth_code(auth_url, self.redirect_uri)
            self.client.auth_provider.authenticate(code, self.redirect_uri,
                                                   self.client_secret)
            with open('.onedrive_session.pkl', 'wb') as pfile:
                pickle.dump(self.client.auth_provider._session, pfile)

        return self.client

    def list_files(self, callback_fn):
        """ list non-directory files """
        def walk_nodes(parentid='root'):
            parent_node = self.client.item(id=parentid)
            for node in t_get(parent_node.children):
                item = node.to_dict()
                item['parentid'] = parentid
                if 'folder' in item and item['folder']['childCount'] > 0:
                    walk_nodes(parentid=item['id'])
                else:
                    callback_fn(item)

        walk_nodes(parentid='root')

    def get_folders(self, callback_fn):
        """ get folders """
        def walk_nodes(parentid='root'):
            parent_node = self.client.item(id=parentid)
            for node in t_get(parent_node.children):
                item = node.to_dict()
                item['parentid'] = parentid
                if 'folder' in item:
                    callback_fn(item)
                    if item['folder']['childCount'] > 0:
                        walk_nodes(parentid=item['id'])
        walk_nodes(parentid='root')

    def download(self, did, exportfile, sha1sum=None):
        """ download using dlink url """
        dirname = os.path.dirname(exportfile)
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        self.client.item(id=did).download(exportfile + '.new')
        if sha1sum:
            from sync_app.util import get_sha1
            sha = get_sha1(exportfile + '.new')
            if sha != sha1sum:
                raise TypeError('%s %s' % (sha, sha1sum))
        os.rename('%s.new' % exportfile, exportfile)
        return True

    def upload(self, fname, parent_id='root'):
        """ upload fname and assign parent_id if provided """
        bname = os.path.basename(fname)
        node = self.client.item(drive="me", id=parent_id)
        return node.children[bname].upload(fname).to_dict()

    def create_directory(self, dname, parent_id='root'):
        """ create directory, assign parent_id if supplied """
        if not parent_id:
            raise ValueError('need to specify parent_id')
        newfolder = Folder()
        newitem = Item()
        newitem.name = dname
        newitem.folder = newfolder

        result = self.client.item(id=parent_id).children.add(newitem).to_dict()
        result['parentid'] = parent_id
        return result

    def delete_file(self, fileid):
        """ delete file by fileid """
        return self.client.item(id=fileid).delete()
<|endoftext|>