<commit_before>from __future__ import absolute_import, division, print_function

import numbers

import torch
from torch.distributions import constraints
from torch.distributions.utils import broadcast_all

from pyro.distributions.torch import Beta, Binomial, Dirichlet, Gamma, Multinomial, Poisson
from pyro.distributions.torch_distribution import TorchDistribution

def _log_beta(x, y):
    return torch.lgamma(x) + torch.lgamma(y) - torch.lgamma(x + y)

def _log_factorial(x):
    return torch.lgamma(1 + x)

class BetaBinomial(TorchDistribution):
    r"""
    Compound distribution comprising of a beta-binomial pair. The probability of
    success (``probs`` for the :class:`~pyro.distributions.Binomial` distribution)
    is unknown and randomly drawn from a :class:`~pyro.distributions.Beta` distribution
    prior to a certain number of Bernoulli trials given by ``total_count``.

    :param float or torch.Tensor concentration1: 1st concentration parameter (alpha) for the
        Beta distribution.
    :param float or torch.Tensor concentration0: 2nd concentration parameter (beta) for the
        Beta distribution.
    :param int or torch.Tensor total_count: number of Bernoulli trials.
    """
    arg_constraints = {'concentration1': constraints.positive, 'concentration0': constraints.positive,
                       'total_count': constraints.nonnegative_integer}
    has_enumerate_support = True
    support = Binomial.support

    def __init__(self, concentration1, concentration0, total_count=1, validate_args=None):
        concentration1, concentration0, total_count = broadcast_all(
            concentration1, concentration0, total_count)
        self._beta = Beta(concentration1, concentration0)
        self.total_count = total_count
        super(BetaBinomial, self).__init__(self._beta._batch_shape, validate_args=validate_args)

    @property
    def concentration1(self):
        return self._beta.concentration1

    @property
    def concentration0(self):
        return self._beta.concentration0

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(BetaBinomial, _instance)
        batch_shape = torch.Size(batch_shape)
        new._beta = self._beta.expand(batch_shape)
        new.total_count = self.total_count.expand_as(new._beta.concentration0)
        super(BetaBinomial, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        probs = self._beta.sample(sample_shape)
        return Binomial(self.total_count, probs).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        log_factorial_n = torch.lgamma(self.total_count + 1)
        log_factorial_k = torch.lgamma(value + 1)
        log_factorial_nmk = torch.lgamma(self.total_count - value + 1)
        return (log_factorial_n - log_factorial_k - log_factorial_nmk +
                _log_beta(value + self.concentration1, self.total_count - value + self.concentration0) -
                _log_beta(self.concentration0, self.concentration1))

    @property
    def mean(self):
        return self._beta.mean * self.total_count

    @property
    def variance(self):
        return self._beta.variance * self.total_count * (self.concentration0 + self.concentration1 + self.total_count)

    def enumerate_support(self, expand=True):
        total_count = int(self.total_count.max())
        if not self.total_count.min() == total_count:
            raise NotImplementedError("Inhomogeneous total count not supported by `enumerate_support`.")
        values = torch.arange(1 + total_count, dtype=self.concentration1.dtype, device=self.concentration1.device)
        values = values.view((-1,) + (1,) * len(self._batch_shape))
        if expand:
            values = values.expand((-1,) + self._batch_shape)
        return values

class DirichletMultinomial(TorchDistribution):
    r"""
    Compound distribution comprising of a dirichlet-multinomial pair. The probability of
    classes (``probs`` for the :class:`~pyro.distributions.Multinomial` distribution)
    is unknown and randomly drawn from a :class:`~pyro.distributions.Dirichlet`
    distribution prior to a certain number of Categorical trials given by
    ``total_count``.

    :param float or torch.Tensor concentration: concentration parameter (alpha) for the
        Dirichlet distribution.
    :param int or torch.Tensor total_count: number of Categorical trials.
    :param bool is_sparse: Whether to assume value is mostly zero when computing
        :meth:`log_prob`, which can speed up computation when data is sparse.
    """
    arg_constraints = {'concentration': constraints.positive, 'total_count': constraints.nonnegative_integer}
    support = Multinomial.support

    def __init__(self, concentration, total_count=1, is_sparse=False, validate_args=None):
        if isinstance(total_count, numbers.Number):
            total_count = concentration.new_tensor(total_count)
        total_count_1 = total_count.unsqueeze(-1)
        concentration, total_count = torch.broadcast_tensors(concentration, total_count_1)
        total_count = total_count_1.squeeze(-1)
        self._dirichlet = Dirichlet(concentration)
        self.total_count = total_count
        self.is_sparse = is_sparse
        super(DirichletMultinomial, self).__init__(
            self._dirichlet._batch_shape, self._dirichlet.event_shape, validate_args=validate_args)

    @property
    def concentration(self):
        return self._dirichlet.concentration

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(DirichletMultinomial, _instance)
        batch_shape = torch.Size(batch_shape)
        new._dirichlet = self._dirichlet.expand(batch_shape)
        new.total_count = self.total_count.expand(batch_shape)
        new.is_sparse = self.is_sparse
        super(DirichletMultinomial, new).__init__(
            new._dirichlet.batch_shape, new._dirichlet.event_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        probs = self._dirichlet.sample(sample_shape)
        total_count = int(self.total_count.max())
        if not self.total_count.min() == total_count:
            raise NotImplementedError("Inhomogeneous total count not supported by `sample`.")
        return Multinomial(total_count, probs).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        n = self.total_count
        alpha = self.concentration
        alpha_sum = self.concentration.sum(-1)
        if self.is_sparse:
            value, alpha = torch.broadcast_tensors(value, alpha)
            elementwise = torch.zeros_like(value)
            mask = (value != 0)
            value = value[mask]
            alpha = alpha[mask]
            elementwise[mask] = torch.lgamma(value + alpha) - _log_factorial(value) - torch.lgamma(alpha)
        else:
            elementwise = torch.lgamma(value + alpha) - _log_factorial(value) - torch.lgamma(alpha)
        return _log_factorial(n) + torch.lgamma(alpha_sum) - torch.lgamma(n + alpha_sum) + elementwise.sum(-1)

    @property
    def mean(self):
        return self._dirichlet.mean * self.total_count.unsqueeze(-1)

    @property
    def variance(self):
        n = self.total_count.unsqueeze(-1)
        alpha = self.concentration
        alpha_sum = self.concentration.sum(-1, keepdim=True)
        alpha_ratio = alpha / alpha_sum
        return n * alpha_ratio * (1 - alpha_ratio) * (n + alpha_sum) / (1 + alpha_sum)

class GammaPoisson(TorchDistribution):
    r"""
    Compound distribution comprising of a gamma-poisson pair, also referred to as
    a gamma-poisson mixture. The ``rate`` parameter for the
    :class:`~pyro.distributions.Poisson` distribution is unknown and randomly
    drawn from a :class:`~pyro.distributions.Gamma` distribution.

    .. note:: This can be treated as an alternate parametrization of the
        :class:`~pyro.distributions.NegativeBinomial` (``total_count``, ``probs``)
        distribution, with `concentration = total_count` and `rate = (1 - probs) / probs`.

    :param float or torch.Tensor concentration: shape parameter (alpha) of the Gamma
        distribution.
    :param float or torch.Tensor rate: rate parameter (beta) for the Gamma
        distribution.
    """

    arg_constraints = {'concentration': constraints.positive, 'rate': constraints.positive}
    support = Poisson.support

    def __init__(self, concentration, rate, validate_args=None):
        concentration, rate = broadcast_all(concentration, rate)
        self._gamma = Gamma(concentration, rate)
        super(GammaPoisson, self).__init__(self._gamma._batch_shape, validate_args=validate_args)

    @property
    def concentration(self):
        return self._gamma.concentration

    @property
    def rate(self):
        return self._gamma.rate

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(GammaPoisson, _instance)
        batch_shape = torch.Size(batch_shape)
        new._gamma = self._gamma.expand(batch_shape)
        super(GammaPoisson, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        rate = self._gamma.sample(sample_shape)
        return Poisson(rate).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        post_value = self.concentration + value
        return -_log_beta(self.concentration, value + 1) - post_value.log() + \
            self.concentration * self.rate.log() - post_value * (1 + self.rate).log()

    @property
    def mean(self):
        return self.concentration / self.rate

    @property
    def variance(self):
        return self.concentration / self.rate.pow(2) * (1 + self.rate)
<commit_msg>Improve sparse DirichletMultinomial.log_prob() (#1743)<commit_after>from __future__ import absolute_import, division, print_function

import numbers

import torch
from torch.distributions import constraints
from torch.distributions.utils import broadcast_all

from pyro.distributions.torch import Beta, Binomial, Dirichlet, Gamma, Multinomial, Poisson
from pyro.distributions.torch_distribution import TorchDistribution

def _log_beta(x, y):
    return torch.lgamma(x) + torch.lgamma(y) - torch.lgamma(x + y)

def _log_beta_1(alpha, value, is_sparse):
    if is_sparse:
        mask = (value != 0)
        value, alpha, mask = torch.broadcast_tensors(value, alpha, mask)
        result = torch.zeros_like(value)
        value = value[mask]
        alpha = alpha[mask]
        result[mask] = torch.lgamma(1 + value) + torch.lgamma(alpha) - torch.lgamma(value + alpha)
        return result
    else:
        return torch.lgamma(1 + value) + torch.lgamma(alpha) - torch.lgamma(value + alpha)

class BetaBinomial(TorchDistribution):
    r"""
    Compound distribution comprising of a beta-binomial pair. The probability of
    success (``probs`` for the :class:`~pyro.distributions.Binomial` distribution)
    is unknown and randomly drawn from a :class:`~pyro.distributions.Beta` distribution
    prior to a certain number of Bernoulli trials given by ``total_count``.

    :param float or torch.Tensor concentration1: 1st concentration parameter (alpha) for the
        Beta distribution.
    :param float or torch.Tensor concentration0: 2nd concentration parameter (beta) for the
        Beta distribution.
    :param int or torch.Tensor total_count: number of Bernoulli trials.
    """
    arg_constraints = {'concentration1': constraints.positive, 'concentration0': constraints.positive,
                       'total_count': constraints.nonnegative_integer}
    has_enumerate_support = True
    support = Binomial.support

    def __init__(self, concentration1, concentration0, total_count=1, validate_args=None):
        concentration1, concentration0, total_count = broadcast_all(
            concentration1, concentration0, total_count)
        self._beta = Beta(concentration1, concentration0)
        self.total_count = total_count
        super(BetaBinomial, self).__init__(self._beta._batch_shape, validate_args=validate_args)

    @property
    def concentration1(self):
        return self._beta.concentration1

    @property
    def concentration0(self):
        return self._beta.concentration0

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(BetaBinomial, _instance)
        batch_shape = torch.Size(batch_shape)
        new._beta = self._beta.expand(batch_shape)
        new.total_count = self.total_count.expand_as(new._beta.concentration0)
        super(BetaBinomial, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        probs = self._beta.sample(sample_shape)
        return Binomial(self.total_count, probs).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        log_factorial_n = torch.lgamma(self.total_count + 1)
        log_factorial_k = torch.lgamma(value + 1)
        log_factorial_nmk = torch.lgamma(self.total_count - value + 1)
        return (log_factorial_n - log_factorial_k - log_factorial_nmk +
                _log_beta(value + self.concentration1, self.total_count - value + self.concentration0) -
                _log_beta(self.concentration0, self.concentration1))

    @property
    def mean(self):
        return self._beta.mean * self.total_count

    @property
    def variance(self):
        return self._beta.variance * self.total_count * (self.concentration0 + self.concentration1 + self.total_count)

    def enumerate_support(self, expand=True):
        total_count = int(self.total_count.max())
        if not self.total_count.min() == total_count:
            raise NotImplementedError("Inhomogeneous total count not supported by `enumerate_support`.")
        values = torch.arange(1 + total_count, dtype=self.concentration1.dtype, device=self.concentration1.device)
        values = values.view((-1,) + (1,) * len(self._batch_shape))
        if expand:
            values = values.expand((-1,) + self._batch_shape)
        return values

class DirichletMultinomial(TorchDistribution):
    r"""
    Compound distribution comprising of a dirichlet-multinomial pair. The probability of
    classes (``probs`` for the :class:`~pyro.distributions.Multinomial` distribution)
    is unknown and randomly drawn from a :class:`~pyro.distributions.Dirichlet`
    distribution prior to a certain number of Categorical trials given by
    ``total_count``.

    :param float or torch.Tensor concentration: concentration parameter (alpha) for the
        Dirichlet distribution.
    :param int or torch.Tensor total_count: number of Categorical trials.
    :param bool is_sparse: Whether to assume value is mostly zero when computing
        :meth:`log_prob`, which can speed up computation when data is sparse.
    """
    arg_constraints = {'concentration': constraints.positive, 'total_count': constraints.nonnegative_integer}
    support = Multinomial.support

    def __init__(self, concentration, total_count=1, is_sparse=False, validate_args=None):
        if isinstance(total_count, numbers.Number):
            total_count = concentration.new_tensor(total_count)
        total_count_1 = total_count.unsqueeze(-1)
        concentration, total_count = torch.broadcast_tensors(concentration, total_count_1)
        total_count = total_count_1.squeeze(-1)
        self._dirichlet = Dirichlet(concentration)
        self.total_count = total_count
        self.is_sparse = is_sparse
        super(DirichletMultinomial, self).__init__(
            self._dirichlet._batch_shape, self._dirichlet.event_shape, validate_args=validate_args)

    @property
    def concentration(self):
        return self._dirichlet.concentration

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(DirichletMultinomial, _instance)
        batch_shape = torch.Size(batch_shape)
        new._dirichlet = self._dirichlet.expand(batch_shape)
        new.total_count = self.total_count.expand(batch_shape)
        new.is_sparse = self.is_sparse
        super(DirichletMultinomial, new).__init__(
            new._dirichlet.batch_shape, new._dirichlet.event_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        probs = self._dirichlet.sample(sample_shape)
        total_count = int(self.total_count.max())
        if not self.total_count.min() == total_count:
            raise NotImplementedError("Inhomogeneous total count not supported by `sample`.")
        return Multinomial(total_count, probs).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        alpha = self.concentration
        return (_log_beta_1(alpha.sum(-1), value.sum(-1), self.is_sparse) -
                _log_beta_1(alpha, value, self.is_sparse).sum(-1))

    @property
    def mean(self):
        return self._dirichlet.mean * self.total_count.unsqueeze(-1)

    @property
    def variance(self):
        n = self.total_count.unsqueeze(-1)
        alpha = self.concentration
        alpha_sum = self.concentration.sum(-1, keepdim=True)
        alpha_ratio = alpha / alpha_sum
        return n * alpha_ratio * (1 - alpha_ratio) * (n + alpha_sum) / (1 + alpha_sum)

class GammaPoisson(TorchDistribution):
    r"""
    Compound distribution comprising of a gamma-poisson pair, also referred to as
    a gamma-poisson mixture. The ``rate`` parameter for the
    :class:`~pyro.distributions.Poisson` distribution is unknown and randomly
    drawn from a :class:`~pyro.distributions.Gamma` distribution.

    .. note:: This can be treated as an alternate parametrization of the
        :class:`~pyro.distributions.NegativeBinomial` (``total_count``, ``probs``)
        distribution, with `concentration = total_count` and `rate = (1 - probs) / probs`.

    :param float or torch.Tensor concentration: shape parameter (alpha) of the Gamma
        distribution.
    :param float or torch.Tensor rate: rate parameter (beta) for the Gamma
        distribution.
    """

    arg_constraints = {'concentration': constraints.positive, 'rate': constraints.positive}
    support = Poisson.support

    def __init__(self, concentration, rate, validate_args=None):
        concentration, rate = broadcast_all(concentration, rate)
        self._gamma = Gamma(concentration, rate)
        super(GammaPoisson, self).__init__(self._gamma._batch_shape, validate_args=validate_args)

    @property
    def concentration(self):
        return self._gamma.concentration

    @property
    def rate(self):
        return self._gamma.rate

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(GammaPoisson, _instance)
        batch_shape = torch.Size(batch_shape)
        new._gamma = self._gamma.expand(batch_shape)
        super(GammaPoisson, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def sample(self, sample_shape=()):
        rate = self._gamma.sample(sample_shape)
        return Poisson(rate).sample()

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        post_value = self.concentration + value
        return -_log_beta(self.concentration, value + 1) - post_value.log() + \
            self.concentration * self.rate.log() - post_value * (1 + self.rate).log()

    @property
    def mean(self):
        return self.concentration / self.rate

    @property
    def variance(self):
        return self.concentration / self.rate.pow(2) * (1 + self.rate)
<|endoftext|>


<commit_before># Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0

import sys

from .messenger import Messenger
from .trace_struct import Trace
from .util import site_is_subsample

def identify_dense_edges(trace):
    """
    Modifies a trace in-place by adding all edges based on the
    `cond_indep_stack` information stored at each site.
    """
    for name, node in trace.nodes.items():
        if site_is_subsample(node):
            continue
        if node["type"] == "sample":
            for past_name, past_node in trace.nodes.items():
                if site_is_subsample(node):
                    continue
                if past_node["type"] == "sample":
                    if past_name == name:
                        break
                    past_node_independent = False
                    for query, target in zip(node["cond_indep_stack"], past_node["cond_indep_stack"]):
                        if query.name == target.name and query.counter != target.counter:
                            past_node_independent = True
                            break
                    if not past_node_independent:
                        trace.add_edge(past_name, name)

class TraceMessenger(Messenger):
    """
    Return a handler that records the inputs and outputs of primitive calls
    and their dependencies.

    Consider the following Pyro program:

        >>> def model(x):
        ...     s = pyro.param("s", torch.tensor(0.5))
        ...     z = pyro.sample("z", dist.Normal(x, s))
        ...     return z ** 2

    We can record its execution using ``trace``
    and use the resulting data structure to compute the log-joint probability
    of all of the sample sites in the execution or extract all parameters.

        >>> trace = pyro.poutine.trace(model).get_trace(0.0)
        >>> logp = trace.log_prob_sum()
        >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]

    :param fn: a stochastic function (callable containing Pyro primitive calls)
    :param graph_type: string that specifies the kind of graph to construct
    :param param_only: if true, only records params and not samples
    :returns: stochastic function decorated with a :class:`~pyro.poutine.trace_messenger.TraceMessenger`
    """

    def __init__(self, graph_type=None, param_only=None):
        """
        :param string graph_type: string that specifies the type of graph
            to construct (currently only "flat" or "dense" supported)
        :param param_only: boolean that specifies whether to record sample sites
        """
        super().__init__()
        if graph_type is None:
            graph_type = "flat"
        if param_only is None:
            param_only = False
        assert graph_type in ("flat", "dense")
        self.graph_type = graph_type
        self.param_only = param_only
        self.trace = Trace(graph_type=self.graph_type)

    def __enter__(self):
        self.trace = Trace(graph_type=self.graph_type)
        return super().__enter__()

    def __exit__(self, *args, **kwargs):
        """
        Adds appropriate edges based on cond_indep_stack information
        upon exiting the context.
        """
        if self.param_only:
            for node in list(self.trace.nodes.values()):
                if node["type"] != "param":
                    self.trace.remove_node(node["name"])
        if self.graph_type == "dense":
            identify_dense_edges(self.trace)
        return super().__exit__(*args, **kwargs)

    def __call__(self, fn):
        """
        TODO docs
        """
        return TraceHandler(self, fn)

    def get_trace(self):
        """
        :returns: data structure
        :rtype: pyro.poutine.Trace

        Helper method for a very common use case.
        Returns a shallow copy of ``self.trace``.
        """
        return self.trace.copy()

    def _reset(self):
        tr = Trace(graph_type=self.graph_type)
        if "_INPUT" in self.trace.nodes:
            tr.add_node("_INPUT",
                        name="_INPUT", type="input",
                        args=self.trace.nodes["_INPUT"]["args"],
                        kwargs=self.trace.nodes["_INPUT"]["kwargs"])
        self.trace = tr
        super()._reset()

    def _pyro_post_sample(self, msg):
        if self.param_only:
            return
        if msg["infer"].get("_do_not_trace"):
            assert msg["infer"].get("is_auxiliary")
            assert not msg["is_observed"]
            return
        self.trace.add_node(msg["name"], **msg.copy())

    def _pyro_post_param(self, msg):
        self.trace.add_node(msg["name"], **msg.copy())

class TraceHandler:
    """
    Execution trace poutine.

    A TraceHandler records the input and output to every Pyro primitive
    and stores them as a site in a Trace().
    This should, in theory, be sufficient information for every inference algorithm
    (along with the implicit computational graph in the Variables?)

    We can also use this for visualization.
    """
    def __init__(self, msngr, fn):
        self.fn = fn
        self.msngr = msngr

    def __call__(self, *args, **kwargs):
        """
        Runs the stochastic function stored in this poutine,
        with additional side effects.

        Resets self.trace to an empty trace,
        installs itself on the global execution stack,
        runs self.fn with the given arguments,
        uninstalls itself from the global execution stack,
        stores the arguments and return value of the function in special sites,
        and returns self.fn's return value
        """
        with self.msngr:
            self.msngr.trace.add_node("_INPUT",
                                      name="_INPUT", type="args",
                                      args=args, kwargs=kwargs)
            try:
                ret = self.fn(*args, **kwargs)
            except (ValueError, RuntimeError):
                exc_type, exc_value, traceback = sys.exc_info()
                shapes = self.msngr.trace.format_shapes()
                raise exc_type(u"{}\n{}".format(exc_value, shapes)).with_traceback(traceback)
            self.msngr.trace.add_node("_RETURN", name="_RETURN", type="return", value=ret)
        return ret

    @property
    def trace(self):
        return self.msngr.trace

    def get_trace(self, *args, **kwargs):
        """
        :returns: data structure
        :rtype: pyro.poutine.Trace

        Helper method for a very common use case.
        Calls this poutine and returns its trace instead of the function's return value.
        """
        self(*args, **kwargs)
        return self.msngr.get_trace()
<commit_msg>Fix reraising logic (#2494)<commit_after># Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0

import sys

from .messenger import Messenger
from .trace_struct import Trace
from .util import site_is_subsample

def identify_dense_edges(trace):
    """
    Modifies a trace in-place by adding all edges based on the
    `cond_indep_stack` information stored at each site.
    """
    for name, node in trace.nodes.items():
        if site_is_subsample(node):
            continue
        if node["type"] == "sample":
            for past_name, past_node in trace.nodes.items():
                if site_is_subsample(node):
                    continue
                if past_node["type"] == "sample":
                    if past_name == name:
                        break
                    past_node_independent = False
                    for query, target in zip(node["cond_indep_stack"], past_node["cond_indep_stack"]):
                        if query.name == target.name and query.counter != target.counter:
                            past_node_independent = True
                            break
                    if not past_node_independent:
                        trace.add_edge(past_name, name)

class TraceMessenger(Messenger):
    """
    Return a handler that records the inputs and outputs of primitive calls
    and their dependencies.

    Consider the following Pyro program:

        >>> def model(x):
        ...     s = pyro.param("s", torch.tensor(0.5))
        ...     z = pyro.sample("z", dist.Normal(x, s))
        ...     return z ** 2

    We can record its execution using ``trace``
    and use the resulting data structure to compute the log-joint probability
    of all of the sample sites in the execution or extract all parameters.

        >>> trace = pyro.poutine.trace(model).get_trace(0.0)
        >>> logp = trace.log_prob_sum()
        >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]

    :param fn: a stochastic function (callable containing Pyro primitive calls)
    :param graph_type: string that specifies the kind of graph to construct
    :param param_only: if true, only records params and not samples
    :returns: stochastic function decorated with a :class:`~pyro.poutine.trace_messenger.TraceMessenger`
    """

    def __init__(self, graph_type=None, param_only=None):
        """
        :param string graph_type: string that specifies the type of graph
            to construct (currently only "flat" or "dense" supported)
        :param param_only: boolean that specifies whether to record sample sites
        """
        super().__init__()
        if graph_type is None:
            graph_type = "flat"
        if param_only is None:
            param_only = False
        assert graph_type in ("flat", "dense")
        self.graph_type = graph_type
        self.param_only = param_only
        self.trace = Trace(graph_type=self.graph_type)

    def __enter__(self):
        self.trace = Trace(graph_type=self.graph_type)
        return super().__enter__()

    def __exit__(self, *args, **kwargs):
        """
        Adds appropriate edges based on cond_indep_stack information
        upon exiting the context.
        """
        if self.param_only:
            for node in list(self.trace.nodes.values()):
                if node["type"] != "param":
                    self.trace.remove_node(node["name"])
        if self.graph_type == "dense":
            identify_dense_edges(self.trace)
        return super().__exit__(*args, **kwargs)

    def __call__(self, fn):
        """
        TODO docs
        """
        return TraceHandler(self, fn)

    def get_trace(self):
        """
        :returns: data structure
        :rtype: pyro.poutine.Trace

        Helper method for a very common use case.
        Returns a shallow copy of ``self.trace``.
        """
        return self.trace.copy()

    def _reset(self):
        tr = Trace(graph_type=self.graph_type)
        if "_INPUT" in self.trace.nodes:
            tr.add_node("_INPUT",
                        name="_INPUT", type="input",
                        args=self.trace.nodes["_INPUT"]["args"],
                        kwargs=self.trace.nodes["_INPUT"]["kwargs"])
        self.trace = tr
        super()._reset()

    def _pyro_post_sample(self, msg):
        if self.param_only:
            return
        if msg["infer"].get("_do_not_trace"):
            assert msg["infer"].get("is_auxiliary")
            assert not msg["is_observed"]
            return
        self.trace.add_node(msg["name"], **msg.copy())

    def _pyro_post_param(self, msg):
        self.trace.add_node(msg["name"], **msg.copy())

class TraceHandler:
    """
    Execution trace poutine.

    A TraceHandler records the input and output to every Pyro primitive
    and stores them as a site in a Trace().
    This should, in theory, be sufficient information for every inference algorithm
    (along with the implicit computational graph in the Variables?)

    We can also use this for visualization.
    """
    def __init__(self, msngr, fn):
        self.fn = fn
        self.msngr = msngr

    def __call__(self, *args, **kwargs):
        """
        Runs the stochastic function stored in this poutine,
        with additional side effects.

        Resets self.trace to an empty trace,
        installs itself on the global execution stack,
        runs self.fn with the given arguments,
        uninstalls itself from the global execution stack,
        stores the arguments and return value of the function in special sites,
        and returns self.fn's return value
        """
        with self.msngr:
            self.msngr.trace.add_node("_INPUT",
                                      name="_INPUT", type="args",
                                      args=args, kwargs=kwargs)
            try:
                ret = self.fn(*args, **kwargs)
            except (ValueError, RuntimeError):
                exc_type, exc_value, traceback = sys.exc_info()
                shapes = self.msngr.trace.format_shapes()
                exc = exc_type(u"{}\n{}".format(exc_value, shapes))
                exc = exc.with_traceback(traceback)
                raise exc from None
            self.msngr.trace.add_node("_RETURN", name="_RETURN", type="return", value=ret)
        return ret

    @property
    def trace(self):
        return self.msngr.trace

    def get_trace(self, *args, **kwargs):
        """
        :returns: data structure
        :rtype: pyro.poutine.Trace

        Helper method for a very common use case.
        Calls this poutine and returns its trace instead of the function's return value.
        """
        self(*args, **kwargs)
        return self.msngr.get_trace()
<|endoftext|>


<commit_before>import datetime as dt
from importlib import reload
import numpy as np
import os
import tempfile

import pysat
from pysat.utils import files as futils
from pysat.tests.ci_test_class import CICleanSetup

class TestBasics():

    temporary_file_list = False

    def setup(self):
        """Runs before every method to create a clean testing setup."""

    def teardown(self):
        """Runs after every method to clean up previous testing."""

    def test_parse_delimited_filename(self):
        """Check ability to parse list of delimited files"""
        # Note: Can be removed if future instrument that uses delimited
        # filenames is added to routine travis end-to-end testing
        fname = ''.join(('test_{year:4d}_{month:2d}_{day:2d}_{hour:2d}',
                         '_{minute:2d}_{second:2d}_{version:2s}_r02.cdf'))
        year = np.ones(6) * 2009
        month = np.ones(6) * 12
        day = np.array([12, 15, 17, 19, 22, 24])
        hour = np.array([8, 10, 6, 18, 3, 23])
        minute = np.array([8, 10, 6, 18, 3, 59])
        second = np.array([58, 11, 26, 2, 18, 59])
        version = np.array(['v1', 'v2', 'r1', 'r3', 'v5', 'a6'])
        file_list = []
        for i in range(6):
            file_list.append(fname.format(year=year[i].astype(int),
                                          month=month[i].astype(int),
                                          day=day[i], hour=hour[i],
                                          minute=minute[i], second=second[i],
                                          version=version[i]))

        file_dict = futils.parse_delimited_filenames(file_list, fname, '_')
        assert np.all(file_dict['year'] == year)
        assert np.all(file_dict['month'] == month)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['hour'] == hour)
        assert np.all(file_dict['minute'] == minute)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['version'] == version)
        assert (file_dict['revision'] is None)
        assert (file_dict['cycle'] is None)

class TestFileDirectoryTranslations(CICleanSetup):

    def setup(self):
        """Runs before every method to create a clean testing setup."""

        # Module is only required for testing installations on TravisCI
        import pysatSpaceWeather

        # Create clean environment on the CI server
        CICleanSetup.setup(self)
        reload(pysat)

        # create temporary directory
        self.tempdir = tempfile.TemporaryDirectory()
        pysat.params['data_dirs'] = [self.tempdir.name]

        # Create several pysat.SpaceWeather instruments and download data.
        # We want to start with a setup that covers general cases a user may
        # have.
        pysat.utils.registry.register_by_module(pysatSpaceWeather.instruments)

        self.insts = []
        self.insts_dates = []
        self.insts_kwargs = []

        # Data by day, ACE SIS data
        self.insts.append(pysat.Instrument('sw', 'ace', tag='historic',
                                           inst_id='sis'))
        test_dates = pysatSpaceWeather.instruments.sw_ace._test_dates
        self.insts_dates.append([test_dates['sis']['historic']] * 2)
        self.insts_kwargs.append({})

        # Data with date mangling, regular F10.7 data, stored monthly
        self.insts.append(pysat.Instrument('sw', 'f107'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates[''][''],
                                 test_dates[''][''] + dt.timedelta(weeks=52)])
        self.insts_kwargs.append({'freq': 'MS'})

        # Data with date mangling, 'all' F10.7 data, single file
        self.insts.append(pysat.Instrument('sw', 'f107', tag='all'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates['']['all'],
                                 test_dates['']['all']])
        self.insts_kwargs.append({})

        # Download data for all instruments
        for inst, dates, kwargs in zip(self.insts, self.insts_dates,
                                       self.insts_kwargs):
            ostr = ' '.join(('Downloading data for', inst.platform,
                             inst.name, inst.tag, inst.inst_id))
            print(ostr)
            inst.download(start=dates[0], stop=dates[1], **kwargs)

    def teardown(self):
        """Runs after every method to clean up previous testing."""

        # Clean environment
        CICleanSetup.teardown(self)

        self.tempdir.cleanup()

    def test_updating_directories(self, capsys):
        """Test directory structure update method"""

        # A variety of options are performed within a single test
        # so that only one round of downloads is performed.

        # New Template
        templ = '{platform}'

        # Convert directories to simpler platform structure, as a test run,
        # with expanded descriptive output.
        futils.update_data_directory_structure(new_template=templ,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check for descriptive output from full_breakdown
        fstr = 'Will move: '
        assert captured.out.find(fstr) >= 0

        # Check how many instruments have no files found. Will be used later.
        index = 0
        orig_num_missing = 0
        while index < len(captured.out):
            index = captured.out.find('No files found.', index)
            if index == -1:
                break
            else:
                index += 1
                orig_num_missing += 1

        # Convert directories to simpler platform structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ' '.join(('All', inst.platform, inst.name, inst.tag,
                             inst.inst_id, 'files moved and accounted for.',
                             '\n'))
            assert captured.out.find(fstr) >= 0

            # Check that directories were removed.
            assert not os.path.isdir(inst.files.data_path)

        # Store newly organized directory format
        pysat.params['directory_format'] = templ

        # Convert directories back to more complex structure
        # First, define new template
        templ = '{platform}/{name}/{tag}/{inst_id}'

        # Update structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ' '.join(('All', inst.platform, inst.name, inst.tag,
                             inst.inst_id, 'files moved and accounted for.',
                             '\n'))
            assert captured.out.find(fstr) >= 0

            # Refresh inst with the old directory template set to get now 'old'
            # path information.
            inst2 = pysat.Instrument(inst.platform, inst.name, tag=inst.tag,
                                     inst_id=inst.inst_id)

            # Check that directories with simpler platform org were NOT removed.
            assert os.path.isdir(inst2.files.data_path)

            # Confirm that the code was aware the directory was not empty.
            fstr = ''.join(('Directory is not empty: ',
                            inst2.files.data_path, '\nEnding cleanup.', '\n'))
            assert captured.out.find(fstr) >= 0

        # Try to update structure again. Files have already moved so
        # no files should be found.
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check for no files output
        index = 0
        num_missing = 0
        while index < len(captured.out):
            index = captured.out.find('No files found.', index)
            if index == -1:
                break
            else:
                index += 1
                num_missing += 1

        # Get difference in number of instruments with no files.
        new_missing = num_missing - orig_num_missing

        # Confirm none of the instruments had files.
        assert new_missing == len(self.insts)

        # Store new format like a typical user would
        pysat.params['directory_format'] = templ
<commit_msg>STY: flake8<commit_after>import datetime as dt
from importlib import reload
import numpy as np
import os
import tempfile

import pysat
from pysat.utils import files as futils
from pysat.tests.ci_test_class import CICleanSetup

class TestBasics():

    temporary_file_list = False

    def setup(self):
        """Runs before every method to create a clean testing setup."""

    def teardown(self):
        """Runs after every method to clean up previous testing."""

    def test_parse_delimited_filename(self):
        """Check ability to parse list of delimited files"""
        # Note: Can be removed if future instrument that uses delimited
        # filenames is added to routine travis end-to-end testing
        fname = ''.join(('test_{year:4d}_{month:2d}_{day:2d}_{hour:2d}',
                         '_{minute:2d}_{second:2d}_{version:2s}_r02.cdf'))
        year = np.ones(6) * 2009
        month = np.ones(6) * 12
        day = np.array([12, 15, 17, 19, 22, 24])
        hour = np.array([8, 10, 6, 18, 3, 23])
        minute = np.array([8, 10, 6, 18, 3, 59])
        second = np.array([58, 11, 26, 2, 18, 59])
        version = np.array(['v1', 'v2', 'r1', 'r3', 'v5', 'a6'])
        file_list = []
        for i in range(6):
            file_list.append(fname.format(year=year[i].astype(int),
                                          month=month[i].astype(int),
                                          day=day[i], hour=hour[i],
                                          minute=minute[i], second=second[i],
                                          version=version[i]))

        file_dict = futils.parse_delimited_filenames(file_list, fname, '_')
        assert np.all(file_dict['year'] == year)
        assert np.all(file_dict['month'] == month)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['hour'] == hour)
        assert np.all(file_dict['minute'] == minute)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['version'] == version)
        assert (file_dict['revision'] is None)
        assert (file_dict['cycle'] is None)

class TestFileDirectoryTranslations(CICleanSetup):

    def setup(self):
        """Runs before every method to create a clean testing setup."""

        # Module is only required for testing installations on TravisCI
        import pysatSpaceWeather

        # Create clean environment on the CI server
        CICleanSetup.setup(self)
        reload(pysat)

        # create temporary directory
        self.tempdir = tempfile.TemporaryDirectory()
        pysat.params['data_dirs'] = [self.tempdir.name]

        # Create several pysat.SpaceWeather instruments and download data.
        # We want to start with a setup that covers general cases a user may
        # have.
        pysat.utils.registry.register_by_module(pysatSpaceWeather.instruments)

        self.insts = []
        self.insts_dates = []
        self.insts_kwargs = []

        # Data by day, ACE SIS data
        self.insts.append(pysat.Instrument('sw', 'ace', tag='historic',
                                           inst_id='sis'))
        test_dates = pysatSpaceWeather.instruments.sw_ace._test_dates
        self.insts_dates.append([test_dates['sis']['historic']] * 2)
        self.insts_kwargs.append({})

        # Data with date mangling, regular F10.7 data, stored monthly
        self.insts.append(pysat.Instrument('sw', 'f107'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates[''][''],
                                 test_dates[''][''] + dt.timedelta(weeks=52)])
        self.insts_kwargs.append({'freq': 'MS'})

        # Data with date mangling, 'all' F10.7 data, single file
        self.insts.append(pysat.Instrument('sw', 'f107', tag='all'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates['']['all'],
                                 test_dates['']['all']])
        self.insts_kwargs.append({})

        # Download data for all instruments
        for inst, dates, kwargs in zip(self.insts, self.insts_dates,
                                       self.insts_kwargs):
            ostr = ' '.join(('Downloading data for', inst.platform,
                             inst.name, inst.tag, inst.inst_id))
            print(ostr)
            inst.download(start=dates[0], stop=dates[1], **kwargs)

    def teardown(self):
        """Runs after every method to clean up previous testing."""

        # Clean environment
        CICleanSetup.teardown(self)

        self.tempdir.cleanup()

    def test_updating_directories(self, capsys):
        """Test directory structure update method"""

        # A variety of options are performed within a single test
        # so that only one round of downloads is performed.

        # New Template
        templ = '{platform}'

        # Convert directories to simpler platform structure, as a test run,
        # with expanded descriptive output.
        futils.update_data_directory_structure(new_template=templ,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check for descriptive output from full_breakdown
        fstr = 'Will move: '
        assert captured.out.find(fstr) >= 0

        # Check how many instruments have no files found. Will be used later.
        index = 0
        orig_num_missing = 0
        while index < len(captured.out):
            index = captured.out.find('No files found.', index)
            if index == -1:
                break
            else:
                index += 1
                orig_num_missing += 1

        # Convert directories to simpler platform structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ' '.join(('All', inst.platform, inst.name, inst.tag,
                             inst.inst_id, 'files moved and accounted for.',
                             '\n'))
            assert captured.out.find(fstr) >= 0

            # Check that directories were removed.
            assert not os.path.isdir(inst.files.data_path)

        # Store newly organized directory format
        pysat.params['directory_format'] = templ

        # Convert directories back to more complex structure
        # First, define new template
        templ = '{platform}/{name}/{tag}/{inst_id}'

        # Update structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ' '.join(('All', inst.platform, inst.name, inst.tag,
                             inst.inst_id, 'files moved and accounted for.',
                             '\n'))
            assert captured.out.find(fstr) >= 0

            # Refresh inst with the old directory template set to get now 'old'
            # path information.
            inst2 = pysat.Instrument(inst.platform, inst.name, tag=inst.tag,
                                     inst_id=inst.inst_id)

            # Check that directories with simpler platform org were NOT removed.
            assert os.path.isdir(inst2.files.data_path)

            # Confirm that the code was aware the directory was not empty.
            fstr = ''.join(('Directory is not empty: ',
                            inst2.files.data_path, '\nEnding cleanup.', '\n'))
            assert captured.out.find(fstr) >= 0

        # Try to update structure again. Files have already moved so
        # no files should be found.
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True,
                                               full_breakdown=True)

        # Capture printouts
        captured = capsys.readouterr()

        # Check for no files output
        index = 0
        num_missing = 0
        while index < len(captured.out):
            index = captured.out.find('No files found.', index)
            if index == -1:
                break
            else:
                index += 1
                num_missing += 1

        # Get difference in number of instruments with no files.
        new_missing = num_missing - orig_num_missing

        # Confirm none of the instruments had files.
        assert new_missing == len(self.insts)

        # Store new format like a typical user would
        pysat.params['directory_format'] = templ
<|endoftext|>


<commit_before>import datetime as dt
import numpy as np
import os

import pysat
from pysat.utils import files as futils
from pysat.tests.travisci_test_class import TravisCICleanSetup

class TestBasics():

    temporary_file_list = False

    def setup(self):
        """Runs before every method to create a clean testing setup."""

    def teardown(self):
        """Runs after every method to clean up previous testing."""

    def test_parse_delimited_filename(self):
        """Check ability to parse list of delimited files"""
        # Note: Can be removed if future instrument that uses delimited
        # filenames is added to routine travis end-to-end testing
        fname = ''.join(('test_{year:4d}_{month:2d}_{day:2d}_{hour:2d}',
                         '_{minute:2d}_{second:2d}_{version:2s}_r02.cdf'))
        year = np.ones(6) * 2009
        month = np.ones(6) * 12
        day = np.array([12, 15, 17, 19, 22, 24])
        hour = np.array([8, 10, 6, 18, 3, 23])
        minute = np.array([8, 10, 6, 18, 3, 59])
        second = np.array([58, 11, 26, 2, 18, 59])
        version = np.array(['v1', 'v2', 'r1', 'r3', 'v5', 'a6'])
        file_list = []
        for i in range(6):
            file_list.append(fname.format(year=year[i].astype(int),
                                          month=month[i].astype(int),
                                          day=day[i], hour=hour[i],
                                          minute=minute[i], second=second[i],
                                          version=version[i]))

        file_dict = futils.parse_delimited_filenames(file_list, fname, '_')
        assert np.all(file_dict['year'] == year)
        assert np.all(file_dict['month'] == month)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['hour'] == hour)
        assert np.all(file_dict['minute'] == minute)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['version'] == version)
        assert (file_dict['revision'] is None)
        assert (file_dict['cycle'] is None)

class TestFileDirectoryTranslations(TravisCICleanSetup):

    def setup(self):
        """Runs before every method to create a clean testing setup."""

        # Module is only required for testing installations on TravisCI
        import pysatSpaceWeather

        # Create clean environment on Travis
        TravisCICleanSetup.setup(self)

        # Create several pysat.SpaceWeather instruments and download data.
        # We want to start with a setup that covers general cases a user may
        # have.
        pysat.utils.registry.register_by_module(pysatSpaceWeather.instruments)

        self.insts = []
        self.insts_dates = []
        self.insts_kwargs = []

        # Data by day, ACE SIS data
        self.insts.append(pysat.Instrument('sw', 'ace', tag='historic',
                                           inst_id='sis'))
        test_dates = pysatSpaceWeather.instruments.sw_ace._test_dates
        self.insts_dates.append([test_dates['sis']['historic']]*2)
        self.insts_kwargs.append({})

        # Data with date mangling, regular F10.7 data, stored monthly
        self.insts.append(pysat.Instrument('sw', 'f107'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates[''][''],
                                 test_dates[''][''] + dt.timedelta(weeks=52)])
        self.insts_kwargs.append({'freq': 'MS'})

        # Data with date mangling, 'all' F10.7 data, single file
        self.insts.append(pysat.Instrument('sw', 'f107', tag='all'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates['']['all'],
                                 test_dates['']['all']])
        self.insts_kwargs.append()

        # Data by 30-days, recent Kp
        self.insts.append(pysat.Instrument('sw', 'kp', tag='recent'))
        test_dates = pysatSpaceWeather.instruments.sw_kp._test_dates
        self.insts_dates.append([test_dates['']['recent']]*2)
        self.insts_kwargs.append({})

        # Data by 3-day, forecast Kp
        self.insts.append(pysat.Instrument('sw', 'kp', tag='forecast'))
        test_dates = pysatSpaceWeather.instruments.sw_kp._test_dates
        self.insts_dates.append([test_dates['']['forecast']]*2)
        self.insts_kwargs.append({})

        # Download data for all instruments
        for inst, dates, kwargs in zip(self.insts, self.insts_dates,
                                       self.insts_kwargs):
            try:
                # TODO: Remove once LASP download issues sorted out
                inst.download(start=dates[0], stop=dates[1], **kwargs)
            except:
                pass

    def teardown(self):
        """Runs after every method to clean up previous testing."""

        # Clean environment on Travis
        TravisCICleanSetup.teardown(self)

        # TODO: Check on potentially removing all inst directories

    def test_updating_directories(self, capsys):

        # New Template
        templ = '{platform}'

        # Convert directories to simpler platform structure, as a test run,
        # with expanded descriptive output.
        futils.update_data_directory_structure(new_template=templ,
                                               full_breakdown=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check for descriptive output from full_breakdown
        fstr = 'Will move: '
        assert captured.find(fstr) >= 0

        # Check how many instruments have no files found. Will be used later.
        index = 0
        orig_num_missing = 0
        while index < len(captured):
            index = captured.find('No files found.', index)
            if index == -1:
                break
            else:
                orig_num_missing += 1

        # Convert directories to simpler platform structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ''.join(('All ', inst.platform, inst.name, inst.tag,
                            inst.inst_id, 'files moved and accounted for.\n'))
            assert captured.find(fstr) >= 0

            # Check that directories were removed.
            assert not os.path.isdir(inst.files.data_path)

        # Store newly organized directory format
        pysat.params['directory_format'] = templ

        # Convert directories back to more complex structure
        # First, define new template
        templ = '{platform}/{name}/{tag}/{inst_id}'

        # Update structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ''.join(('All ', inst.platform, inst.name, inst.tag,
                            inst.inst_id, 'files moved and accounted for.\n'))
            assert captured.find(fstr) >= 0

            # Refresh inst with the old directory template set to get now 'old'
            # path information.
            inst2 = pysat.Instrument(inst.platform, inst.name, tag=inst.tag,
                                     inst_id=inst.inst_id)

            # Check that directories with simpler platform org were NOT removed.
            assert os.path.isdir(inst2.files.data_path)

            # Confirm that the code was aware the directory was not empty.
            fstr = ''.join(('Directory is not empty: ', inst2.files.data_path,
                            ' Ending cleanup.'))
            assert captured.find(fstr) >= 0

        # Try to update structure again. Files have already moved so
        # no files should be found.
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check for no files output
        index = 0
        num_missing = 0
        while index < len(captured):
            index = captured.find('No files found.', index)
            if index == -1:
                break
            else:
                num_missing += 1

        # Get difference in number of instruments with no files.
        new_missing = orig_num_missing - num_missing

        # Confirm none of the instruments had files.
        assert new_missing == len(self.insts)

        # Store new format like a typical user would
        pysat.params['directory_format'] = templ
<commit_msg>STY: docstring<commit_after>import datetime as dt
import numpy as np
import os

import pysat
from pysat.utils import files as futils
from pysat.tests.travisci_test_class import TravisCICleanSetup

class TestBasics():

    temporary_file_list = False

    def setup(self):
        """Runs before every method to create a clean testing setup."""

    def teardown(self):
        """Runs after every method to clean up previous testing."""

    def test_parse_delimited_filename(self):
        """Check ability to parse list of delimited files"""
        # Note: Can be removed if future instrument that uses delimited
        # filenames is added to routine travis end-to-end testing
        fname = ''.join(('test_{year:4d}_{month:2d}_{day:2d}_{hour:2d}',
                         '_{minute:2d}_{second:2d}_{version:2s}_r02.cdf'))
        year = np.ones(6) * 2009
        month = np.ones(6) * 12
        day = np.array([12, 15, 17, 19, 22, 24])
        hour = np.array([8, 10, 6, 18, 3, 23])
        minute = np.array([8, 10, 6, 18, 3, 59])
        second = np.array([58, 11, 26, 2, 18, 59])
        version = np.array(['v1', 'v2', 'r1', 'r3', 'v5', 'a6'])
        file_list = []
        for i in range(6):
            file_list.append(fname.format(year=year[i].astype(int),
                                          month=month[i].astype(int),
                                          day=day[i], hour=hour[i],
                                          minute=minute[i], second=second[i],
                                          version=version[i]))

        file_dict = futils.parse_delimited_filenames(file_list, fname, '_')
        assert np.all(file_dict['year'] == year)
        assert np.all(file_dict['month'] == month)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['hour'] == hour)
        assert np.all(file_dict['minute'] == minute)
        assert np.all(file_dict['day'] == day)
        assert np.all(file_dict['version'] == version)
        assert (file_dict['revision'] is None)
        assert (file_dict['cycle'] is None)

class TestFileDirectoryTranslations(TravisCICleanSetup):

    def setup(self):
        """Runs before every method to create a clean testing setup."""

        # Module is only required for testing installations on TravisCI
        import pysatSpaceWeather

        # Create clean environment on Travis
        TravisCICleanSetup.setup(self)

        # Create several pysat.SpaceWeather instruments and download data.
        # We want to start with a setup that covers general cases a user may
        # have.
        pysat.utils.registry.register_by_module(pysatSpaceWeather.instruments)

        self.insts = []
        self.insts_dates = []
        self.insts_kwargs = []

        # Data by day, ACE SIS data
        self.insts.append(pysat.Instrument('sw', 'ace', tag='historic',
                                           inst_id='sis'))
        test_dates = pysatSpaceWeather.instruments.sw_ace._test_dates
        self.insts_dates.append([test_dates['sis']['historic']]*2)
        self.insts_kwargs.append({})

        # Data with date mangling, regular F10.7 data, stored monthly
        self.insts.append(pysat.Instrument('sw', 'f107'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates[''][''],
                                 test_dates[''][''] + dt.timedelta(weeks=52)])
        self.insts_kwargs.append({'freq': 'MS'})

        # Data with date mangling, 'all' F10.7 data, single file
        self.insts.append(pysat.Instrument('sw', 'f107', tag='all'))
        test_dates = pysatSpaceWeather.instruments.sw_f107._test_dates
        self.insts_dates.append([test_dates['']['all'],
                                 test_dates['']['all']])
        self.insts_kwargs.append()

        # Data by 30-days, recent Kp
        self.insts.append(pysat.Instrument('sw', 'kp', tag='recent'))
        test_dates = pysatSpaceWeather.instruments.sw_kp._test_dates
        self.insts_dates.append([test_dates['']['recent']]*2)
        self.insts_kwargs.append({})

        # Data by 3-day, forecast Kp
        self.insts.append(pysat.Instrument('sw', 'kp', tag='forecast'))
        test_dates = pysatSpaceWeather.instruments.sw_kp._test_dates
        self.insts_dates.append([test_dates['']['forecast']]*2)
        self.insts_kwargs.append({})

        # Download data for all instruments
        for inst, dates, kwargs in zip(self.insts, self.insts_dates,
                                       self.insts_kwargs):
            try:
                # TODO: Remove once LASP download issues sorted out
                inst.download(start=dates[0], stop=dates[1], **kwargs)
            except:
                pass

    def teardown(self):
        """Runs after every method to clean up previous testing."""

        # Clean environment on Travis
        TravisCICleanSetup.teardown(self)

        # TODO: Check on potentially removing all inst directories

    def test_updating_directories(self, capsys):
        """Test directory structure update method"""

        # A variety of options are performed within a single test
        # so that only one round of downloads is performed.

        # New Template
        templ = '{platform}'

        # Convert directories to simpler platform structure, as a test run,
        # with expanded descriptive output.
        futils.update_data_directory_structure(new_template=templ,
                                               full_breakdown=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check for descriptive output from full_breakdown
        fstr = 'Will move: '
        assert captured.find(fstr) >= 0

        # Check how many instruments have no files found. Will be used later.
        index = 0
        orig_num_missing = 0
        while index < len(captured):
            index = captured.find('No files found.', index)
            if index == -1:
                break
            else:
                orig_num_missing += 1

        # Convert directories to simpler platform structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ''.join(('All ', inst.platform, inst.name, inst.tag,
                            inst.inst_id, 'files moved and accounted for.\n'))
            assert captured.find(fstr) >= 0

            # Check that directories were removed.
            assert not os.path.isdir(inst.files.data_path)

        # Store newly organized directory format
        pysat.params['directory_format'] = templ

        # Convert directories back to more complex structure
        # First, define new template
        templ = '{platform}/{name}/{tag}/{inst_id}'

        # Update structure
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check if we did things correctly. Look for correct output strings.
        for inst in self.insts:
            # Check for all files moved.
            fstr = ''.join(('All ', inst.platform, inst.name, inst.tag,
                            inst.inst_id, 'files moved and accounted for.\n'))
            assert captured.find(fstr) >= 0

            # Refresh inst with the old directory template set to get now 'old'
            # path information.
            inst2 = pysat.Instrument(inst.platform, inst.name, tag=inst.tag,
                                     inst_id=inst.inst_id)

            # Check that directories with simpler platform org were NOT removed.
            assert os.path.isdir(inst2.files.data_path)

            # Confirm that the code was aware the directory was not empty.
            fstr = ''.join(('Directory is not empty: ', inst2.files.data_path,
                            ' Ending cleanup.'))
            assert captured.find(fstr) >= 0

        # Try to update structure again. Files have already moved so
        # no files should be found.
        futils.update_data_directory_structure(new_template=templ,
                                               test_run=False,
                                               remove_empty_dirs=True)

        # Capture printouts
        captured, err = capsys.readouterr()

        # Check for no files output
        index = 0
        num_missing = 0
        while index < len(captured):
            index = captured.find('No files found.', index)
            if index == -1:
                break
            else:
                num_missing += 1

        # Get difference in number of instruments with no files.
        new_missing = orig_num_missing - num_missing

        # Confirm none of the instruments had files.
        assert new_missing == len(self.insts)

        # Store new format like a typical user would
        pysat.params['directory_format'] = templ
<|endoftext|>


<commit_before>"""Astroid hooks for the Python 2 GObject introspection bindings.

Helps with understanding everything imported from 'gi.repository'
"""

import inspect
import itertools
import sys
import re
import warnings

from astroid import MANAGER, AstroidBuildingError, nodes
from astroid.builder import AstroidBuilder

_inspected_modules = {}

_identifier_re = r'^[A-Za-z_]\w*$'

def _gi_build_stub(parent):
    """
    Inspect the passed module recursively and build stubs for functions,
    classes, etc.
    """
    classes = {}
    functions = {}
    constants = {}
    methods = {}
    for name in dir(parent):
        if name.startswith("__"):
            continue

        # Check if this is a valid name in python
        if not re.match(_identifier_re, name):
            continue

        try:
            obj = getattr(parent, name)
        except:
            continue

        if inspect.isclass(obj):
            classes[name] = obj
        elif (inspect.isfunction(obj) or
              inspect.isbuiltin(obj)):
            functions[name] = obj
        elif (inspect.ismethod(obj) or
              inspect.ismethoddescriptor(obj)):
            methods[name] = obj
        elif (str(obj).startswith("<flags") or
              str(obj).startswith("<enum ") or
              str(obj).startswith("<GType ") or
              inspect.isdatadescriptor(obj)):
            constants[name] = 0
        elif isinstance(obj, (int, str)):
            constants[name] = obj
        elif callable(obj):
            # Fall back to a function for anything callable
            functions[name] = obj
        else:
            # Assume everything else is some manner of constant
            constants[name] = 0

    ret = ""

    if constants:
        ret += "# %s contants\n\n" % parent.__name__
    for name in sorted(constants):
        if name[0].isdigit():
            # GDK has some busted constant names like
            # Gdk.EventType.2BUTTON_PRESS
            continue

        val = constants[name]

        strval = str(val)
        if isinstance(val, str):
            strval = '"%s"' % str(val).replace("\\", "\\\\")
        ret += "%s = %s\n" % (name, strval)

    if ret:
        ret += "\n\n"
    if functions:
        ret += "# %s functions\n\n" % parent.__name__
    for name in sorted(functions):
        ret += "def %s(*args, **kwargs):\n" % name
        ret += "    pass\n"

    if ret:
        ret += "\n\n"
    if methods:
        ret += "# %s methods\n\n" % parent.__name__
    for name in sorted(methods):
        ret += "def %s(self, *args, **kwargs):\n" % name
        ret += "    pass\n"

    if ret:
        ret += "\n\n"
    if classes:
        ret += "# %s classes\n\n" % parent.__name__
    for name in sorted(classes):
        ret += "class %s(object):\n" % name

        classret = _gi_build_stub(classes[name])
        if not classret:
            classret = "pass\n"

        for line in classret.splitlines():
            ret += "    " + line + "\n"
        ret += "\n"

    return ret

def _import_gi_module(modname):
    # we only consider gi.repository submodules
    if not modname.startswith('gi.repository.'):
        raise AstroidBuildingError(modname=modname)
    # build astroid representation unless we already tried so
    if modname not in _inspected_modules:
        modnames = [modname]
        optional_modnames = []

        # GLib and GObject may have some special case handling
        # in pygobject that we need to cope with. However at
        # least as of pygobject3-3.13.91 the _glib module doesn't
        # exist anymore, so if treat these modules as optional.
        if modname == 'gi.repository.GLib':
            optional_modnames.append('gi._glib')
        elif modname == 'gi.repository.GObject':
            optional_modnames.append('gi._gobject')

        try:
            modcode = ''
            for m in itertools.chain(modnames, optional_modnames):
                try:
                    __import__(m)
                    with warnings.catch_warnings():
                        # Just inspecting the code can raise gi deprecation
                        # warnings, so ignore them.
                        try:
                            from gi import PyGIDeprecationWarning
                            warnings.simplefilter("ignore", PyGIDeprecationWarning)
                        except Exception:
                            pass

                        modcode += _gi_build_stub(sys.modules[m])
                except ImportError:
                    if m not in optional_modnames:
                        raise
        except ImportError:
            astng = _inspected_modules[modname] = None
        else:
            astng = AstroidBuilder(MANAGER).string_build(modcode, modname)
            _inspected_modules[modname] = astng
    else:
        astng = _inspected_modules[modname]
    if astng is None:
        raise AstroidBuildingError(modname=modname)
    return astng

def _looks_like_require_version(node):
    # Return whether this looks like a call to gi.require_version(<name>, <version>)
    # Only accept function calls with two constant arguments
    if len(node.args) != 2:
        return False

    if not all(isinstance(arg, nodes.Const) for arg in node.args):
        return False

    func = node.func
    if isinstance(func, nodes.Attribute):
        if func.attrname != 'require_version':
            return False
        if isinstance(func.expr, nodes.Name) and func.expr.name == 'gi':
            return True

        return False

    if isinstance(func, nodes.Name):
        return func.name == 'require_version'

    return False

def _register_require_version(node):
    # Load the gi.require_version locally
    try:
        import gi
        gi.require_version(node.args[0].value, node.args[1].value)
    except Exception:
        pass

    return node

MANAGER.register_failed_import_hook(_import_gi_module)
MANAGER.register_transform(nodes.Call, _register_require_version, _looks_like_require_version)
<commit_msg>Ignore PyGIWarning<commit_after>"""Astroid hooks for the Python 2 GObject introspection bindings.

Helps with understanding everything imported from 'gi.repository'
"""

import inspect
import itertools
import sys
import re
import warnings

from astroid import MANAGER, AstroidBuildingError, nodes
from astroid.builder import AstroidBuilder

_inspected_modules = {}

_identifier_re = r'^[A-Za-z_]\w*$'

def _gi_build_stub(parent):
    """
    Inspect the passed module recursively and build stubs for functions,
    classes, etc.
    """
    classes = {}
    functions = {}
    constants = {}
    methods = {}
    for name in dir(parent):
        if name.startswith("__"):
            continue

        # Check if this is a valid name in python
        if not re.match(_identifier_re, name):
            continue

        try:
            obj = getattr(parent, name)
        except:
            continue

        if inspect.isclass(obj):
            classes[name] = obj
        elif (inspect.isfunction(obj) or
              inspect.isbuiltin(obj)):
            functions[name] = obj
        elif (inspect.ismethod(obj) or
              inspect.ismethoddescriptor(obj)):
            methods[name] = obj
        elif (str(obj).startswith("<flags") or
              str(obj).startswith("<enum ") or
              str(obj).startswith("<GType ") or
              inspect.isdatadescriptor(obj)):
            constants[name] = 0
        elif isinstance(obj, (int, str)):
            constants[name] = obj
        elif callable(obj):
            # Fall back to a function for anything callable
            functions[name] = obj
        else:
            # Assume everything else is some manner of constant
            constants[name] = 0

    ret = ""

    if constants:
        ret += "# %s contants\n\n" % parent.__name__
    for name in sorted(constants):
        if name[0].isdigit():
            # GDK has some busted constant names like
            # Gdk.EventType.2BUTTON_PRESS
            continue

        val = constants[name]

        strval = str(val)
        if isinstance(val, str):
            strval = '"%s"' % str(val).replace("\\", "\\\\")
        ret += "%s = %s\n" % (name, strval)

    if ret:
        ret += "\n\n"
    if functions:
        ret += "# %s functions\n\n" % parent.__name__
    for name in sorted(functions):
        ret += "def %s(*args, **kwargs):\n" % name
        ret += "    pass\n"

    if ret:
        ret += "\n\n"
    if methods:
        ret += "# %s methods\n\n" % parent.__name__
    for name in sorted(methods):
        ret += "def %s(self, *args, **kwargs):\n" % name
        ret += "    pass\n"

    if ret:
        ret += "\n\n"
    if classes:
        ret += "# %s classes\n\n" % parent.__name__
    for name in sorted(classes):
        ret += "class %s(object):\n" % name

        classret = _gi_build_stub(classes[name])
        if not classret:
            classret = "pass\n"

        for line in classret.splitlines():
            ret += "    " + line + "\n"
        ret += "\n"

    return ret

def _import_gi_module(modname):
    # we only consider gi.repository submodules
    if not modname.startswith('gi.repository.'):
        raise AstroidBuildingError(modname=modname)
    # build astroid representation unless we already tried so
    if modname not in _inspected_modules:
        modnames = [modname]
        optional_modnames = []

        # GLib and GObject may have some special case handling
        # in pygobject that we need to cope with. However at
        # least as of pygobject3-3.13.91 the _glib module doesn't
        # exist anymore, so if treat these modules as optional.
        if modname == 'gi.repository.GLib':
            optional_modnames.append('gi._glib')
        elif modname == 'gi.repository.GObject':
            optional_modnames.append('gi._gobject')

        try:
            modcode = ''
            for m in itertools.chain(modnames, optional_modnames):
                try:
                    with warnings.catch_warnings():
                        # Just inspecting the code can raise gi deprecation
                        # warnings, so ignore them.
                        try:
                            from gi import PyGIDeprecationWarning, PyGIWarning
                            warnings.simplefilter("ignore", PyGIDeprecationWarning)
                            warnings.simplefilter("ignore", PyGIWarning)
                        except Exception:
                            pass

                        __import__(m)
                        modcode += _gi_build_stub(sys.modules[m])
                except ImportError:
                    if m not in optional_modnames:
                        raise
        except ImportError:
            astng = _inspected_modules[modname] = None
        else:
            astng = AstroidBuilder(MANAGER).string_build(modcode, modname)
            _inspected_modules[modname] = astng
    else:
        astng = _inspected_modules[modname]
    if astng is None:
        raise AstroidBuildingError(modname=modname)
    return astng

def _looks_like_require_version(node):
    # Return whether this looks like a call to gi.require_version(<name>, <version>)
    # Only accept function calls with two constant arguments
    if len(node.args) != 2:
        return False

    if not all(isinstance(arg, nodes.Const) for arg in node.args):
        return False

    func = node.func
    if isinstance(func, nodes.Attribute):
        if func.attrname != 'require_version':
            return False
        if isinstance(func.expr, nodes.Name) and func.expr.name == 'gi':
            return True

        return False

    if isinstance(func, nodes.Name):
        return func.name == 'require_version'

    return False

def _register_require_version(node):
    # Load the gi.require_version locally
    try:
        import gi
        gi.require_version(node.args[0].value, node.args[1].value)
    except Exception:
        pass

    return node

MANAGER.register_failed_import_hook(_import_gi_module)
MANAGER.register_transform(nodes.Call, _register_require_version, _looks_like_require_version)
<|endoftext|>


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

"""Functionality for interacting with the orders api"""

import logging
from urllib.parse import urljoin

from .http import PlanetSession
from . import auth
from . import models

PLANET_BASE_URL =  'https://api.planet.com/'
BASE_URL = PLANET_BASE_URL + 'compute/ops/'
STATS_PATH = 'stats/orders/v2/'
ORDERS_PATH = 'orders/v2/'

LOGGER = logging.getLogger(__name__)

class OrdersClientException(Exception):
    """Exceptions thrown by Orders Client"""
    pass

class OrdersClient(object):
    """High-level access to Planet's orders API.

    Basic Usage::
      >>>
      >> from planet.api.orders_client import OrdersClient
      >> cl = OrdersClient('api_key')
      >> order = cl.get_order('order_id')
      <models.Order>
    """
    def __init__(self, api_key=None, base_url=BASE_URL):
        '''
        :param str api_key: (opt) API key to use.
            Defaults to environment variable or stored authentication data.
        :param str base_url: (opt) The base URL to use. Defaults to production
            orders API base url.
        '''
        api_key = api_key or auth.find_api_key()
        self.auth = api_key and auth.APIKey(api_key)

        self._base_url = base_url
        if not self._base_url.endswith('/'):
            self._base_url += '/'

    def _check_order_id(self, oid):
        if not oid or not isinstance(oid, str):
            msg = 'Order id ({}) is invalid.'.format(oid)
            raise OrdersClientException(msg)

    def _orders_url(self):
        return self._base_url + ORDERS_PATH

    def _stats_url(self):
        return self._base_url + STATS_PATH

    def _order_url(self, order_id):
        self._check_order_id
        return self._orders_url() + order_id

    def _request(self, url, method, body_type, data=None):
        '''Prepare an order API request.

        :param url str: location to make request
        :returns: :py:Class:`planet.api.models.Request`
        '''
        return models.Request(url, self.auth, body_type=body_type,
                              method=method, data=data)

    def _get(self, url, body_type):
        request = self._request(url, 'GET', body_type)
        return self._do_request(request)

    def _put(self, url, body_type):
        request = self._request(url, 'PUT', body_type)
        return self._do_request(request)

    def _post(self, url, body_type, data):
        request = self._request(url, 'POST', body_type, data)
        return self._do_request(request)

    def _do_request(self, request):
        with PlanetSession() as sess:
            body = sess.request(request).body
        return body

    def list_orders(self):
        '''Get all order requests.

        :returns: :py:Class:`planet.api.models.Orders`
        :raises planet.api.exceptions.APIException: On API error.
        '''
        # url = self.base_url
        #
        # request = self._request(url, 'GET', models.Orders)
        #
        # with PlanetSession() as sess:
        #     orders = sess.request(request).body
        #
        # return orders
        raise NotImplementedError

    def create_order(self, order_request):
        '''Create an order request.

        :param dict order_request: order request details
        :returns str: The ID of the order
        '''
        if not isinstance(order_request, models.OrderDetails):
            order_request = models.OrderDetails(order_request)

        url = self._orders_url()

        order = self._post(url, models.Order, order_request.data)
        return order.id

    def get_order(self, order_id):
        '''Get order details by Order ID.

        :param order_id str: The ID of the order
        :returns: :py:Class:`planet.api.models.Order`
        :raises planet.api.exceptions.APIException: On API error.
        '''
        url = self._order_url(order_id)
        order = self._get(url, models.Order)
        return order

    def cancel_order(self, order_id):
        '''Cancel a queued order.

        :param order_id str: The ID of the order
        :returns :py:Class:`planet.api.models.Order`: Cancelled order details.
        :raises planet.api.exceptions.APIException: On API error.
        '''
        url = self._order_url(order_id)
        order = self._put(url, models.Order)
        return order

    def cancel_orders(self, order_ids):
        '''Cancel queued orders in bulk.
        '''
        raise NotImplementedError

    def aggregated_order_stats(self):
        '''Get aggregated counts of active orders.

        :returns dict: aggregated order counts
        :raises planet.api.exceptions.APIException: On API error.

        '''
        url = self._stats_url()
        res = self._get(url, models.JSON)
        return res.data

    def download_order(self, order_id, wait=False):
        '''Download ordered asset

        :param order_id str: The ID of the order
        '''
        raise NotImplementedError

# class BaseClient(object):
#     def _url(self, path):
#         if path.startswith('http'):
#             url = path
#         else:
#             url = self.base_url + path
#         return url
#
#     def _request(self, path, body_type=models.JSON, params=None, auth=None):
#         return models.Request(self._url(path), auth or self.auth, params,
#                               body_type)
#
#     def _get(self, path, body_type=models.JSON, params=None, callback=None):
#         # convert any JSON objects to text explicitly
#         for k, v in (params or {}).items():
#             if isinstance(v, dict):
#                 params[k] = json.dumps(v)
#
#         request = self._request(path, body_type, params)
#         response = self.dispatcher.response(request)
#         if callback:
#             response.get_body_async(callback)
#         return response
<commit_msg>pyflakes fixups<commit_after># Copyright 2020 Planet Labs, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
# import logging

"""Functionality for interacting with the orders api"""

import logging

from .http import PlanetSession
from . import auth
from . import models

PLANET_BASE_URL = 'https://api.planet.com/'
BASE_URL = PLANET_BASE_URL + 'compute/ops/'
STATS_PATH = 'stats/orders/v2/'
ORDERS_PATH = 'orders/v2/'

LOGGER = logging.getLogger(__name__)

class OrdersClientException(Exception):
    """Exceptions thrown by Orders Client"""
    pass

class OrdersClient(object):
    """High-level access to Planet's orders API.

    Basic Usage::
      >>>
      >> from planet.api.orders_client import OrdersClient
      >> cl = OrdersClient('api_key')
      >> order = cl.get_order('order_id')
      <models.Order>
    """
    def __init__(self, api_key=None, base_url=BASE_URL):
        '''
        :param str api_key: (opt) API key to use.
            Defaults to environment variable or stored authentication data.
        :param str base_url: (opt) The base URL to use. Defaults to production
            orders API base url.
        '''
        api_key = api_key or auth.find_api_key()
        self.auth = api_key and auth.APIKey(api_key)

        self._base_url = base_url
        if not self._base_url.endswith('/'):
            self._base_url += '/'

    def _check_order_id(self, oid):
        if not oid or not isinstance(oid, str):
            msg = 'Order id ({}) is invalid.'.format(oid)
            raise OrdersClientException(msg)

    def _orders_url(self):
        return self._base_url + ORDERS_PATH

    def _stats_url(self):
        return self._base_url + STATS_PATH

    def _order_url(self, order_id):
        self._check_order_id
        return self._orders_url() + order_id

    def _request(self, url, method, body_type, data=None):
        '''Prepare an order API request.

        :param url str: location to make request
        :returns: :py:Class:`planet.api.models.Request`
        '''
        return models.Request(url, self.auth, body_type=body_type,
                              method=method, data=data)

    def _get(self, url, body_type):
        request = self._request(url, 'GET', body_type)
        return self._do_request(request)

    def _put(self, url, body_type):
        request = self._request(url, 'PUT', body_type)
        return self._do_request(request)

    def _post(self, url, body_type, data):
        request = self._request(url, 'POST', body_type, data)
        return self._do_request(request)

    def _do_request(self, request):
        with PlanetSession() as sess:
            body = sess.request(request).body
        return body

    def list_orders(self):
        '''Get all order requests.

        :returns: :py:Class:`planet.api.models.Orders`
        :raises planet.api.exceptions.APIException: On API error.
        '''
        # url = self.base_url
        #
        # request = self._request(url, 'GET', models.Orders)
        #
        # with PlanetSession() as sess:
        #     orders = sess.request(request).body
        #
        # return orders
        raise NotImplementedError

    def create_order(self, order_request):
        '''Create an order request.

        :param dict order_request: order request details
        :returns str: The ID of the order
        '''
        if not isinstance(order_request, models.OrderDetails):
            order_request = models.OrderDetails(order_request)

        url = self._orders_url()

        order = self._post(url, models.Order, order_request.data)
        return order.id

    def get_order(self, order_id):
        '''Get order details by Order ID.

        :param order_id str: The ID of the order
        :returns: :py:Class:`planet.api.models.Order`
        :raises planet.api.exceptions.APIException: On API error.
        '''
        url = self._order_url(order_id)
        order = self._get(url, models.Order)
        return order

    def cancel_order(self, order_id):
        '''Cancel a queued order.

        :param order_id str: The ID of the order
        :returns :py:Class:`planet.api.models.Order`: Cancelled order details.
        :raises planet.api.exceptions.APIException: On API error.
        '''
        url = self._order_url(order_id)
        order = self._put(url, models.Order)
        return order

    def cancel_orders(self, order_ids):
        '''Cancel queued orders in bulk.
        '''
        raise NotImplementedError

    def aggregated_order_stats(self):
        '''Get aggregated counts of active orders.

        :returns dict: aggregated order counts
        :raises planet.api.exceptions.APIException: On API error.

        '''
        url = self._stats_url()
        res = self._get(url, models.JSON)
        return res.data

    def download_order(self, order_id, wait=False):
        '''Download ordered asset

        :param order_id str: The ID of the order
        '''
        raise NotImplementedError

# class BaseClient(object):
#     def _url(self, path):
#         if path.startswith('http'):
#             url = path
#         else:
#             url = self.base_url + path
#         return url
#
#     def _request(self, path, body_type=models.JSON, params=None, auth=None):
#         return models.Request(self._url(path), auth or self.auth, params,
#                               body_type)
#
#     def _get(self, path, body_type=models.JSON, params=None, callback=None):
#         # convert any JSON objects to text explicitly
#         for k, v in (params or {}).items():
#             if isinstance(v, dict):
#                 params[k] = json.dumps(v)
#
#         request = self._request(path, body_type, params)
#         response = self.dispatcher.response(request)
#         if callback:
#             response.get_body_async(callback)
#         return response
<|endoftext|>