<commit_before>from typing import Iterable

from fireant.utils import immutable
from pypika.terms import (
    Case,
    NullValue,
    ValueWrapper,
)
from .base import SlicerElement
from .exceptions import QueryException
from .filters import (
    AntiPatternFilter,
    BooleanFilter,
    ContainsFilter,
    ExcludesFilter,
    PatternFilter,
    RangeFilter,
)
from .intervals import (
    NumericInterval,
    daily,
)

class Dimension(SlicerElement):
    """
    The `Dimension` class represents a dimension in the `Slicer` object.

    :param alias:
        A unique identifier used to identify the metric when writing slicer queries. This value must be unique over
        the metrics in the slicer.

    :param definition:
        A pypika expression which is used to select the value when building SQL queries.
    """

    def __init__(self, key, label=None, definition=None, display_definition=None):
        super(Dimension, self).__init__(key, label, definition, display_definition)
        self.is_rollup = False

    @immutable
    def rollup(self):
        """
        Configures this dimension and all subsequent dimensions in a slicer query to be rolled up to provide the totals.
        This will include an extra value for each pair of dimensions labeled `Totals`. which will include the totals for
        the group.
        """
        self.is_rollup = True

    def __repr__(self):
        return "slicer.dimensions.{}".format(self.key)

class BooleanDimension(Dimension):
    """
    This is a dimension that represents a boolean true/false value.  The expression should always result in a boolean
    value.
    """

    def __init__(self, key, label=None, definition=None):
        super(BooleanDimension, self).__init__(key,
                                               label,
                                               definition)

    def is_(self, value: bool):
        """
        Creates a filter to filter a slicer query.

        :param value:
            True or False
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is True or False.
        """
        return BooleanFilter(self.definition, value)

class PatternFilterableMixin:
    definition = None
    pattern_definition_attribute = 'definition'

    def like(self, pattern, *patterns):
        """
        Creates a filter to filter a slicer query.

        :param pattern:
            A pattern to match against the dimension's display definition.  This pattern is used in the SQL query as
            the `LIKE` expression.
        :param patterns:
            Additional patterns. This is the same as the pattern argument. The function signature is intended to
            syntactically require at least one pattern.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension's display definition
            matches the pattern.
        """
        return PatternFilter(getattr(self, self.pattern_definition_attribute), pattern, *patterns)

    def not_like(self, pattern, *patterns):
        """
        Creates a filter to filter a slicer query.

        :param pattern:
            A pattern to match against the dimension's display definition.  This pattern is used in the SQL query as
            the `NOT LIKE` expression.
        :param patterns:
            Additional patterns. This is the same as the pattern argument. The function signature is intended to
            syntactically require at least one pattern.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension's display definition
            matches the pattern.
        """
        return AntiPatternFilter(getattr(self, self.pattern_definition_attribute), pattern, *patterns)

class CategoricalDimension(PatternFilterableMixin, Dimension):
    """
    This is a dimension that represents an enum-like database field, with a finite list of options to chose from. It
    provides support for configuring a display value for each of the possible values.
    """

    def __init__(self, key, label=None, definition=None, display_values=()):
        super(CategoricalDimension, self).__init__(key,
                                                   label,
                                                   definition)
        self.display_values = dict(display_values)

    def isin(self, values: Iterable):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is one of a set of
            values. Opposite of #notin.
        """
        return ContainsFilter(self.definition, values)

    def notin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is *not* one of a set of
            values. Opposite of #isin.
        """
        return ExcludesFilter(self.definition, values)

class _UniqueDimensionBase(PatternFilterableMixin, Dimension):
    def isin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is one of a set of
            values. Opposite of #notin.
        """
        return ContainsFilter(self.definition, values)

    def notin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is *not* one of a set of
            values. Opposite of #isin.
        """
        return ExcludesFilter(self.definition, values)

class UniqueDimension(_UniqueDimensionBase):
    """
    This is a dimension that represents a field in a database which is a unique identifier, such as a primary/foreign
    key. It provides support for a display value field which is selected and used in the results.
    """
    def __init__(self, key, label=None, definition=None, display_definition=None):
        super(UniqueDimension, self).__init__(key,
                                              label,
                                              definition,
                                              display_definition)
        if display_definition is not None:
            self.display = DisplayDimension(self)

    @property
    def has_display_field(self):
        return hasattr(self, 'display')

    def like(self, pattern, *patterns):
        if not self.has_display_field:
            raise QueryException('No value set for display_definition.')
        return self.display.like(pattern, *patterns)

    def not_like(self, pattern, *patterns):
        if not self.has_display_field:
            raise QueryException('No value set for display_definition.')
        return self.display.not_like(pattern, *patterns)

class DisplayDimension(_UniqueDimensionBase):
    """
    WRITEME
    """

    def __init__(self, dimension):
        super(DisplayDimension, self).__init__('{}_display'.format(dimension.key),
                                               dimension.label,
                                               dimension.display_definition)

class ContinuousDimension(Dimension):
    """
    This is a dimension that represents a field in the database which is a continuous value, such as a decimal, integer,
    or date/time. It requires the use of an interval which is the window over the values.
    """

    def __init__(self, key, label=None, definition=None, default_interval=NumericInterval(1, 0)):
        super(ContinuousDimension, self).__init__(key,
                                                  label,
                                                  definition)
        self.interval = default_interval

class DatetimeDimension(ContinuousDimension):
    """
    A subclass of ContinuousDimension which reflects a date/time data type. Intervals are replaced with time intervals
    such as daily, weekly, annually, etc.  A reference can be used to show a comparison over time such as
    week-over-week or month-over-month.
    """

    def __init__(self, key, label=None, definition=None, default_interval=daily):
        super(DatetimeDimension, self).__init__(key,
                                                label,
                                                definition,
                                                default_interval=default_interval)

    @immutable
    def __call__(self, interval):
        """
        When calling a datetime dimension an interval can be supplied:

        .. code-block:: python

            from fireant import weekly

            my_slicer.dimensions.date # Daily interval used as default
            my_slicer.dimensions.date(weekly) # Daily interval used as default

        :param interval:
            An interval to use with the dimension.  See `fireant.intervals`.
        :return:
            A copy of the dimension with the interval set.
        """
        self.interval = interval

    def between(self, start, stop):
        """
        Creates a filter to filter a slicer query.

        :param start:
            The start time of the filter. This is the beginning of the window for which results should be included.
        :param stop:
            The stop time of the filter. This is the end of the window for which results should be included.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is between the values
            start and stop.
        """
        return RangeFilter(self.definition, start, stop)

class TotalsDimension(Dimension):
    def __init__(self, dimension):
        totals_definition = NullValue()
        display_definition = totals_definition \
            if dimension.has_display_field \
            else None

        super(TotalsDimension, self).__init__(dimension.key,
                                              dimension.label,
                                              totals_definition,
                                              display_definition)
<commit_msg>Added a SQLLite dialect Created a method to override the default ValueWrapper implementation Created a SQLLiteValueWrapper which serializes booleans as 1s and 0s<commit_after>from typing import Iterable

from fireant.utils import immutable
from pypika.terms import NullValue
from .base import SlicerElement
from .exceptions import QueryException
from .filters import (
    AntiPatternFilter,
    BooleanFilter,
    ContainsFilter,
    ExcludesFilter,
    PatternFilter,
    RangeFilter,
)
from .intervals import (
    NumericInterval,
    daily,
)

class Dimension(SlicerElement):
    """
    The `Dimension` class represents a dimension in the `Slicer` object.

    :param alias:
        A unique identifier used to identify the metric when writing slicer queries. This value must be unique over
        the metrics in the slicer.

    :param definition:
        A pypika expression which is used to select the value when building SQL queries.
    """

    def __init__(self, key, label=None, definition=None, display_definition=None):
        super(Dimension, self).__init__(key, label, definition, display_definition)
        self.is_rollup = False

    @immutable
    def rollup(self):
        """
        Configures this dimension and all subsequent dimensions in a slicer query to be rolled up to provide the totals.
        This will include an extra value for each pair of dimensions labeled `Totals`. which will include the totals for
        the group.
        """
        self.is_rollup = True

    def __repr__(self):
        return "slicer.dimensions.{}".format(self.key)

class BooleanDimension(Dimension):
    """
    This is a dimension that represents a boolean true/false value.  The expression should always result in a boolean
    value.
    """

    def __init__(self, key, label=None, definition=None):
        super(BooleanDimension, self).__init__(key,
                                               label,
                                               definition)

    def is_(self, value: bool):
        """
        Creates a filter to filter a slicer query.

        :param value:
            True or False
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is True or False.
        """
        return BooleanFilter(self.definition, value)

class PatternFilterableMixin:
    definition = None
    pattern_definition_attribute = 'definition'

    def like(self, pattern, *patterns):
        """
        Creates a filter to filter a slicer query.

        :param pattern:
            A pattern to match against the dimension's display definition.  This pattern is used in the SQL query as
            the `LIKE` expression.
        :param patterns:
            Additional patterns. This is the same as the pattern argument. The function signature is intended to
            syntactically require at least one pattern.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension's display definition
            matches the pattern.
        """
        return PatternFilter(getattr(self, self.pattern_definition_attribute), pattern, *patterns)

    def not_like(self, pattern, *patterns):
        """
        Creates a filter to filter a slicer query.

        :param pattern:
            A pattern to match against the dimension's display definition.  This pattern is used in the SQL query as
            the `NOT LIKE` expression.
        :param patterns:
            Additional patterns. This is the same as the pattern argument. The function signature is intended to
            syntactically require at least one pattern.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension's display definition
            matches the pattern.
        """
        return AntiPatternFilter(getattr(self, self.pattern_definition_attribute), pattern, *patterns)

class CategoricalDimension(PatternFilterableMixin, Dimension):
    """
    This is a dimension that represents an enum-like database field, with a finite list of options to chose from. It
    provides support for configuring a display value for each of the possible values.
    """

    def __init__(self, key, label=None, definition=None, display_values=()):
        super(CategoricalDimension, self).__init__(key,
                                                   label,
                                                   definition)
        self.display_values = dict(display_values)

    def isin(self, values: Iterable):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is one of a set of
            values. Opposite of #notin.
        """
        return ContainsFilter(self.definition, values)

    def notin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is *not* one of a set of
            values. Opposite of #isin.
        """
        return ExcludesFilter(self.definition, values)

class _UniqueDimensionBase(PatternFilterableMixin, Dimension):
    def isin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is one of a set of
            values. Opposite of #notin.
        """
        return ContainsFilter(self.definition, values)

    def notin(self, values):
        """
        Creates a filter to filter a slicer query.

        :param values:
            An iterable of value to constrain the slicer query results by.

        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is *not* one of a set of
            values. Opposite of #isin.
        """
        return ExcludesFilter(self.definition, values)

class UniqueDimension(_UniqueDimensionBase):
    """
    This is a dimension that represents a field in a database which is a unique identifier, such as a primary/foreign
    key. It provides support for a display value field which is selected and used in the results.
    """
    def __init__(self, key, label=None, definition=None, display_definition=None):
        super(UniqueDimension, self).__init__(key,
                                              label,
                                              definition,
                                              display_definition)
        if display_definition is not None:
            self.display = DisplayDimension(self)

    @property
    def has_display_field(self):
        return hasattr(self, 'display')

    def like(self, pattern, *patterns):
        if not self.has_display_field:
            raise QueryException('No value set for display_definition.')
        return self.display.like(pattern, *patterns)

    def not_like(self, pattern, *patterns):
        if not self.has_display_field:
            raise QueryException('No value set for display_definition.')
        return self.display.not_like(pattern, *patterns)

class DisplayDimension(_UniqueDimensionBase):
    """
    WRITEME
    """

    def __init__(self, dimension):
        super(DisplayDimension, self).__init__('{}_display'.format(dimension.key),
                                               dimension.label,
                                               dimension.display_definition)

class ContinuousDimension(Dimension):
    """
    This is a dimension that represents a field in the database which is a continuous value, such as a decimal, integer,
    or date/time. It requires the use of an interval which is the window over the values.
    """

    def __init__(self, key, label=None, definition=None, default_interval=NumericInterval(1, 0)):
        super(ContinuousDimension, self).__init__(key,
                                                  label,
                                                  definition)
        self.interval = default_interval

class DatetimeDimension(ContinuousDimension):
    """
    A subclass of ContinuousDimension which reflects a date/time data type. Intervals are replaced with time intervals
    such as daily, weekly, annually, etc.  A reference can be used to show a comparison over time such as
    week-over-week or month-over-month.
    """

    def __init__(self, key, label=None, definition=None, default_interval=daily):
        super(DatetimeDimension, self).__init__(key,
                                                label,
                                                definition,
                                                default_interval=default_interval)

    @immutable
    def __call__(self, interval):
        """
        When calling a datetime dimension an interval can be supplied:

        .. code-block:: python

            from fireant import weekly

            my_slicer.dimensions.date # Daily interval used as default
            my_slicer.dimensions.date(weekly) # Daily interval used as default

        :param interval:
            An interval to use with the dimension.  See `fireant.intervals`.
        :return:
            A copy of the dimension with the interval set.
        """
        self.interval = interval

    def between(self, start, stop):
        """
        Creates a filter to filter a slicer query.

        :param start:
            The start time of the filter. This is the beginning of the window for which results should be included.
        :param stop:
            The stop time of the filter. This is the end of the window for which results should be included.
        :return:
            A slicer query filter used to filter a slicer query to results where this dimension is between the values
            start and stop.
        """
        return RangeFilter(self.definition, start, stop)

class TotalsDimension(Dimension):
    def __init__(self, dimension):
        totals_definition = NullValue()
        display_definition = totals_definition \
            if dimension.has_display_field \
            else None

        super(TotalsDimension, self).__init__(dimension.key,
                                              dimension.label,
                                              totals_definition,
                                              display_definition)
<|endoftext|>


<commit_before>import urllib2
import requests
import pdb
import time

class PyBingException(Exception):
    pass

class PyBingSearch(object):

    QUERY_URL = 'https://api.datamarket.azure.com/Bing/Search/Web' \
                 + '?Query={}&$top={}&$skip={}&$format={}'

    def __init__(self, api_key, safe=False):
        self.api_key = api_key
        self.safe = safe

    def search(self, query, limit=50, offset=0, format='json'):
        return self._search(query, limit, offset, format)

    def search_all(self, query, limit=50, format='json'):
        results = self._search(query, limit, 0, format)
        while results.total > len(results) and len(results) < limit:
            max = limit - len(results)
            more_results = self._search(query, max, len(results), format)
            results += more_results
        return results

    def _search(self, query, limit, offset, format):
        '''
        Returns a list of result objects, with the url for the next page bing search url.
        '''
        url = self.QUERY_URL.format(urllib2.quote("'{}'".format(query)), limit, offset, format)
        r = requests.get(url, auth=("", self.api_key))
        try:
            json_results = r.json()
        except ValueError as vE:
            if not self.safe:
                raise PyBingException("Request returned with code %s, error msg: %s" % (r.status_code, r.text))
            else:
                print "[ERROR] Request returned with code %s, error msg: %s. \nContinuing in 5 seconds." % (r.status_code, r.text)
                time.slee(5)
        try:
            next_link = json_results['d']['__next']
        except KeyError as kE:
            if not self.safe:
                raise PyBingException("Couldn't extract next_link: KeyError: %s" % kE)
            else:
                print "Couldn't extract next_link: KeyError: %s" % kE
            next_link = ''
        return [Result(single_result_json) for single_result_json in json_results['d']['results']], next_link

class Result(object):
    '''
    The class represents a SINGLE search result.
    Each result will come with the following:

    #For the actual results#
    title: title of the result
    url: the url of the result
    description: description for the result
    id: bing id for the page

    #Meta info#:
    meta.uri: the search uri for bing
    meta.type: for the most part WebResult
    '''

    class _Meta(object):
        '''
        Holds the meta info for the result.
        '''
        def __init__(self, meta):
            self.type = meta['type']
            self.uri = meta['uri']

    def __init__(self, result):
        self.url = result['Url']
        self.title = result['Title']
        self.description = result['Description']
        self.id = result['ID']

        self.meta = self._Meta(result['__metadata'])
<commit_msg>fixing merge + adding comments<commit_after>import urllib2
import requests
import pdb
import time

class PyBingException(Exception):
    pass

class PyBingSearch(object):

    QUERY_URL = 'https://api.datamarket.azure.com/Bing/Search/Web' \
                 + '?Query={}&$top={}&$skip={}&$format={}'

    def __init__(self, api_key, safe=False):
        self.api_key = api_key
        self.safe = safe

    def search(self, query, limit=50, offset=0, format='json'):
        ''' Returns the result list, and also the uri for next page (returned_list, next_uri) '''
        return self._search(query, limit, offset, format)

    def search_all(self, query, limit=50, format='json'):
        ''' Returns a single list containing up to 'limit' Result objects'''
        results, next_link = self._search(query, limit, 0, format)
        while next_link and len(results) < limit:
            max = limit - len(results)
            more_results, next_link = self._search(query, max, len(results), format)
            if not more_results:
                break
            results += more_results
        return results

    def _search(self, query, limit, offset, format):
        '''
        Returns a list of result objects, with the url for the next page bing search url.
        '''
        url = self.QUERY_URL.format(urllib2.quote("'{}'".format(query)), limit, offset, format)
        r = requests.get(url, auth=("", self.api_key))
        try:
            json_results = r.json()
        except ValueError as vE:
            if not self.safe:
                raise PyBingException("Request returned with code %s, error msg: %s" % (r.status_code, r.text))
            else:
                print "[ERROR] Request returned with code %s, error msg: %s. \nContinuing in 5 seconds." % (r.status_code, r.text)
                time.sleep(5)
        try:
            next_link = json_results['d']['__next']
        except KeyError as kE:
            if not self.safe:
                raise PyBingException("Couldn't extract next_link: KeyError: %s" % kE)
            else:
                print "Couldn't extract next_link: KeyError: %s" % kE
                time.sleep(3)
            next_link = ''
        return [Result(single_result_json) for single_result_json in json_results['d']['results']], next_link

class Result(object):
    '''
    The class represents a SINGLE search result.
    Each result will come with the following:

    #For the actual results#
    title: title of the result
    url: the url of the result
    description: description for the result
    id: bing id for the page

    #Meta info#:
    meta.uri: the search uri for bing
    meta.type: for the most part WebResult
    '''

    class _Meta(object):
        '''
        Holds the meta info for the result.
        '''
        def __init__(self, meta):
            self.type = meta['type']
            self.uri = meta['uri']

    def __init__(self, result):
        self.url = result['Url']
        self.title = result['Title']
        self.description = result['Description']
        self.id = result['ID']

        self.meta = self._Meta(result['__metadata'])
<|endoftext|>


<commit_before># Copyright 2021 The Flax Authors.
#
# 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.

"""Utilities we could consider upstreaming to Jax.
"""

import collections
from collections.abc import Iterable  # pylint: disable=g-importing-member
import warnings

import jax
from jax import lax
from jax import linear_util as lu
from jax.config import config
from jax.interpreters import partial_eval as pe
from jax.interpreters import xla
import jax.lib.xla_bridge as xb
import jax.numpy as jnp
import numpy as np

def _replicate(x, devices=None):
  x = jax.numpy.asarray(x)
  if devices is None:
    # match the default device assignments used in pmap:
    # for single-host, that's the XLA default device assignment
    # for multi-host, it's the order of jax.local_devices()
    if jax.process_count() == 1:
      devices = [d for d in xb.get_backend().get_default_device_assignment(
          jax.device_count()) if d.process_index == jax.process_index()]
    else:
      devices = jax.local_devices()
  if hasattr(jax, "device_put_sharded"):  # jax >= 0.2.0
    return jax.device_put_sharded(len(devices) * [x], devices)
  else:
    aval = jax.ShapedArray((len(devices),) + x.shape, x.dtype)
    buffers = [xla.device_put(x, device=d) for d in devices]
    return jax.pxla.ShardedDeviceArray(aval, buffers)

def replicate(tree, devices=None):
  """Replicates arrays to multiple devices.

  Args:
    tree: a pytree containing the arrays that should be replicated.
    devices: the devices the data is replicated to
      (default: `jax.local_devices()`).
  Returns:
    A new pytree containing the replicated arrays.
  """
  return jax.tree_map(lambda x: _replicate(x, devices), tree)

def unreplicate(tree):
  """Returns a single instance of a replicated array."""
  return jax.tree_map(lambda x: x[0], tree)

def pmean(xs, axis_name):
  warnings.warn('use jax.lax.pmean instead',
                DeprecationWarning)
  return lax.pmean(xs, axis_name)

def partial_eval_by_shape(fn, input_spec, *args, **kwargs):
  """Lazily evaluate a function by using the shapes of the inputs.

  This function is similar to `jax.eval_shape` with the key difference that
  function outputs that can be computed without a concrete value of the
  inputs are returned as is instead of only the shape. See for example
  `module.init_by_shape` where this functionality is used to initialize a
  model without using input data lr computation.

  Args:
    fn: the function to be lazily evaluated.
    input_spec: an iterable of shapes or (shape, dtype) tuples specifying the
      shape and type of the inputs. If unspecified the dtype is float32.
    *args: other arguments passed to the module's apply function
    **kwargs: keyword arguments passed to the module's apply function
  Returns:
    A pair consisting of the model output and an instance of Model
  """
  # output cannot be returned in lazy_create because jax.eval_shape will only
  # return the shape and dtype.
  # TODO(mattjj,jheek): use a public JAX API
  f = lambda *inputs: fn(*inputs, *args, **kwargs)
  input_structs = [_parse_spec(spec) for spec in input_spec]
  inputs_flat, in_tree = jax.tree_flatten(input_structs)
  f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
  in_pvals = [pe.PartialVal.unknown(jax.ShapedArray(x.shape, x.dtype))
              for x in inputs_flat]
  _, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)
  out_flat = [const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)
              for pv, const in out_pvals]
  return jax.tree_unflatten(out_tree(), out_flat)

def _parse_spec(spec):
  """Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct."""
  spec = tuple(spec)
  if len(spec) == 2 and isinstance(spec[0], Iterable):
    return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])
  else:
    return jax.ShapeDtypeStruct(spec, jnp.float32)

def prefetch_to_device(iterator, size, devices=None):
  """"Shard and prefetch batches on device.

  This utility takes an iterator and returns a new iterator which fills an on
  device prefetch buffer. Eager prefetching can improve the performance of
  training loops significantly by overlapping compute and data transfer.

  This utility is mostly useful for GPUs, for TPUs it should not be necessary.

  Args:
    iterator: an iterator that yields a pytree of ndarrays where the first
      dimension is sharded across devices.
    size: the size of the prefetch buffer.
    devices: the list of devices to which the arrays should be prefetched.
  Yields:
    The original items from the iterator where each ndarray is now a sharded to
    the specified devices.
  """
  queue = collections.deque()
  if devices is None:
    devices = jax.local_devices()
  def _prefetch(xs):
    if hasattr(jax, "device_put_sharded"):  # jax>=0.2.0
      return jax.device_put_sharded(list(xs), devices)
    else:
      aval = jax.xla.abstractify(xs)
      assert xs.shape[0] == len(devices), (
          "The first dimension of the iterator's ndarrays is not "
          "equal to the number of devices.")
      buffers = [xla.device_put(x, devices[i])
                 for i, x in enumerate(xs)]
      return jax.pxla.ShardedDeviceArray(aval, buffers)
  try:
    while len(queue) < size:
      queue.append(jax.tree_map(_prefetch, next(iterator)))
  except StopIteration:
    pass

  while True:
    try:
      xs = queue.popleft()
    except IndexError:
      return
    try:
      queue.append(jax.tree_map(_prefetch, next(iterator)))
    except StopIteration:
      pass
    yield xs

def _scan_nd(body_fn, init, xs, n=1):
  """Utility for performing an n-dimensional `lax.scan`.

  The n-d scan is simply recursive call of 1-d scan.
  Args:
    body_fn: the body of the loop of type (c, x) -> (c, y).
    init: initial value for the carry.
    xs: a pytree of tensors to scan over.
    n: number of dimensions to scan over (default: 1)
  Returns:
    A tuple of the final carry and the values returned by the body.
  """
  if n == 1:
    return lax.scan(body_fn, init, xs)
  else:
    def scan_body(c, x):
      return _scan_nd(body_fn, c, x, n=n-1)
    return lax.scan(scan_body, init, xs)

def _invert_perm(perm):
  perm_inv = [0] * len(perm)
  for i, j in enumerate(perm):
    perm_inv[j] = i
  return tuple(perm_inv)

def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):
  """utility for doing a scan along arbitrary dimensions.

  see `lax.scan` for details on how the scan operation works.
  Args:
    body_fn: the body of the loop of type (c, x) -> (c, y).
    init: initial value for the carry.
    xs: a pytree of tensors to scan over.
    axis: the axis to scan over.
    keepdims: keep the dimensions that are scanned over.
  Returns:
    A tuple of the final carry and the values returned by the body.
  """
  if not isinstance(axis, Iterable):
    axis = (axis,)

  def transpose_in(x):
    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
    return x.transpose(perm)
  def transpose_out(x):
    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
    return x.transpose(_invert_perm(perm))

  def body_wrapper(c, xs):
    if keepdims:
      xs = jax.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)
      xs = jax.tree_map(transpose_out, xs)
    c, ys = body_fn(c, xs)
    if keepdims:
      ys = jax.tree_map(transpose_in, ys)
      ys = jax.tree_map(lambda x: x.reshape(x.shape[len(axis):]), ys)
    return c, ys

  xs = jax.tree_map(transpose_in, xs)
  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))
  ys = jax.tree_map(transpose_out, ys)
  return c, ys
<commit_msg>Implement jax_utils.replicate using jax.device_put_replicated<commit_after># Copyright 2021 The Flax Authors.
#
# 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.

"""Utilities we could consider upstreaming to Jax.
"""

import collections
from collections.abc import Iterable  # pylint: disable=g-importing-member
import warnings

import jax
from jax import lax
from jax import linear_util as lu
from jax.config import config
from jax.interpreters import partial_eval as pe
from jax.interpreters import xla
import jax.lib.xla_bridge as xb
import jax.numpy as jnp
import numpy as np

def replicate(tree, devices=None):
  """Replicates arrays to multiple devices.

  Args:
    tree: a pytree containing the arrays that should be replicated.
    devices: the devices the data is replicated to
      (default: `jax.local_devices()`).
  Returns:
    A new pytree containing the replicated arrays.
  """
  if devices is None:
    # match the default device assignments used in pmap:
    # for single-host, that's the XLA default device assignment
    # for multi-host, it's the order of jax.local_devices()
    if jax.process_count() == 1:
      devices = [d for d in xb.get_backend().get_default_device_assignment(
          jax.device_count()) if d.process_index == jax.process_index()]
    else:
      devices = jax.local_devices()

  return jax.device_put_replicated(tree, devices)

def unreplicate(tree):
  """Returns a single instance of a replicated array."""
  return jax.tree_map(lambda x: x[0], tree)

def pmean(xs, axis_name):
  warnings.warn('use jax.lax.pmean instead',
                DeprecationWarning)
  return lax.pmean(xs, axis_name)

def partial_eval_by_shape(fn, input_spec, *args, **kwargs):
  """Lazily evaluate a function by using the shapes of the inputs.

  This function is similar to `jax.eval_shape` with the key difference that
  function outputs that can be computed without a concrete value of the
  inputs are returned as is instead of only the shape. See for example
  `module.init_by_shape` where this functionality is used to initialize a
  model without using input data lr computation.

  Args:
    fn: the function to be lazily evaluated.
    input_spec: an iterable of shapes or (shape, dtype) tuples specifying the
      shape and type of the inputs. If unspecified the dtype is float32.
    *args: other arguments passed to the module's apply function
    **kwargs: keyword arguments passed to the module's apply function
  Returns:
    A pair consisting of the model output and an instance of Model
  """
  # output cannot be returned in lazy_create because jax.eval_shape will only
  # return the shape and dtype.
  # TODO(mattjj,jheek): use a public JAX API
  f = lambda *inputs: fn(*inputs, *args, **kwargs)
  input_structs = [_parse_spec(spec) for spec in input_spec]
  inputs_flat, in_tree = jax.tree_flatten(input_structs)
  f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
  in_pvals = [pe.PartialVal.unknown(jax.ShapedArray(x.shape, x.dtype))
              for x in inputs_flat]
  _, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)
  out_flat = [const if pv is None else jax.ShapeDtypeStruct(pv.shape, pv.dtype)
              for pv, const in out_pvals]
  return jax.tree_unflatten(out_tree(), out_flat)

def _parse_spec(spec):
  """Parse an input spec of the form (shape, dtype) or shape into a jax.ShapeDtypeStruct."""
  spec = tuple(spec)
  if len(spec) == 2 and isinstance(spec[0], Iterable):
    return jax.ShapeDtypeStruct(tuple(spec[0]), spec[1])
  else:
    return jax.ShapeDtypeStruct(spec, jnp.float32)

def prefetch_to_device(iterator, size, devices=None):
  """"Shard and prefetch batches on device.

  This utility takes an iterator and returns a new iterator which fills an on
  device prefetch buffer. Eager prefetching can improve the performance of
  training loops significantly by overlapping compute and data transfer.

  This utility is mostly useful for GPUs, for TPUs it should not be necessary.

  Args:
    iterator: an iterator that yields a pytree of ndarrays where the first
      dimension is sharded across devices.
    size: the size of the prefetch buffer.
    devices: the list of devices to which the arrays should be prefetched.
  Yields:
    The original items from the iterator where each ndarray is now a sharded to
    the specified devices.
  """
  queue = collections.deque()
  if devices is None:
    devices = jax.local_devices()
  def _prefetch(xs):
    if hasattr(jax, "device_put_sharded"):  # jax>=0.2.0
      return jax.device_put_sharded(list(xs), devices)
    else:
      aval = jax.xla.abstractify(xs)
      assert xs.shape[0] == len(devices), (
          "The first dimension of the iterator's ndarrays is not "
          "equal to the number of devices.")
      buffers = [xla.device_put(x, devices[i])
                 for i, x in enumerate(xs)]
      return jax.pxla.ShardedDeviceArray(aval, buffers)
  try:
    while len(queue) < size:
      queue.append(jax.tree_map(_prefetch, next(iterator)))
  except StopIteration:
    pass

  while True:
    try:
      xs = queue.popleft()
    except IndexError:
      return
    try:
      queue.append(jax.tree_map(_prefetch, next(iterator)))
    except StopIteration:
      pass
    yield xs

def _scan_nd(body_fn, init, xs, n=1):
  """Utility for performing an n-dimensional `lax.scan`.

  The n-d scan is simply recursive call of 1-d scan.
  Args:
    body_fn: the body of the loop of type (c, x) -> (c, y).
    init: initial value for the carry.
    xs: a pytree of tensors to scan over.
    n: number of dimensions to scan over (default: 1)
  Returns:
    A tuple of the final carry and the values returned by the body.
  """
  if n == 1:
    return lax.scan(body_fn, init, xs)
  else:
    def scan_body(c, x):
      return _scan_nd(body_fn, c, x, n=n-1)
    return lax.scan(scan_body, init, xs)

def _invert_perm(perm):
  perm_inv = [0] * len(perm)
  for i, j in enumerate(perm):
    perm_inv[j] = i
  return tuple(perm_inv)

def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):
  """utility for doing a scan along arbitrary dimensions.

  see `lax.scan` for details on how the scan operation works.
  Args:
    body_fn: the body of the loop of type (c, x) -> (c, y).
    init: initial value for the carry.
    xs: a pytree of tensors to scan over.
    axis: the axis to scan over.
    keepdims: keep the dimensions that are scanned over.
  Returns:
    A tuple of the final carry and the values returned by the body.
  """
  if not isinstance(axis, Iterable):
    axis = (axis,)

  def transpose_in(x):
    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
    return x.transpose(perm)
  def transpose_out(x):
    perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
    return x.transpose(_invert_perm(perm))

  def body_wrapper(c, xs):
    if keepdims:
      xs = jax.tree_map(lambda x: x.reshape((1,) * len(axis) + x.shape), xs)
      xs = jax.tree_map(transpose_out, xs)
    c, ys = body_fn(c, xs)
    if keepdims:
      ys = jax.tree_map(transpose_in, ys)
      ys = jax.tree_map(lambda x: x.reshape(x.shape[len(axis):]), ys)
    return c, ys

  xs = jax.tree_map(transpose_in, xs)
  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))
  ys = jax.tree_map(transpose_out, ys)
  return c, ys
<|endoftext|>


<commit_before>from splitwise import Splitwise
from splitwise.comment import Comment
import unittest
try:
    from unittest.mock import patch
except ImportError:  # Python 2
    from mock import patch

@patch('splitwise.Splitwise._Splitwise__makeRequest')
class CreateCommentTestCase(unittest.TestCase):

    def setUp(self):
        self.sObj = Splitwise('consumerkey', 'consumersecret')

    def test_createComment_success(self, mockMakeRequest):
        mockMakeRequest.return_value = '{"comment": {"id": 211110000, "content": "Test for create comment", "comment_type": "User", "relation_type": "ExpenseComment", "relation_id": 982430660, "created_at": "2020-08-09T16:14:52Z", "deleted_at": null, "user": {"id": 784241, "first_name": "ruks", "last_name": null, "picture": {"medium": "https://s3.amazonaws.com/splitwise/uploads/user/default_avatars/avatar-ruby47-100px.png"}}}, "errors":{}}'  # noqa: E501
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'
        comment, errors = self.sObj.createComment(data)
        mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                           "POST",
                                           data)
        self.assertEqual(comment.getId(), 211110000)
        self.assertEqual(comment.getContent(), 'Test for create comment')
        self.assertEqual(comment.getCommentType(), 'User')
        self.assertEqual(comment.getRelationType(), 'ExpenseComment')
        self.assertEqual(comment.getRelationId(), 982430660)
        self.assertEqual(comment.getCreatedAt(), '2020-08-09T16:14:52Z')
        self.assertEqual(comment.getDeletedAt(), None)
        user = comment.getCommentedUser()
        self.assertEqual(user.getFirstName(), 'ruks')
        self.assertEqual(user.getLastName(), None)

    def test_createComment_error(self, mockMakeRequest):
        mockMakeRequest.return_value = '{"comment":{},"errors":{"base":["An unknown error occurred. Please try again or contact support@splitwise.com if you experience repeated issues. Sorry for the trouble!"]}}'  # noqa: E501
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'

        comment, errors = self.sObj.createComment(data)
        mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                           "POST",
                                           data)
        self.assertEqual(errors.getErrors(), {
            'base': [
                'An unknown error occurred. Please try again or contact support@splitwise.com if you experience repeated issues. \
Sorry for the trouble!'
            ]})

    def test_createComment_exception(self, mockMakeRequest):
        mockMakeRequest.side_effect = Exception(
            "Invalid response %s. Please check your consumer key and secret." % 404)
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'
        with self.assertRaises(Exception):
            comment, errors = self.sObj.createComment(data)
            mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                               "POST",
                                               data)
<commit_msg>Rectified lint errors createComment<commit_after>from splitwise import Splitwise
import unittest
try:
    from unittest.mock import patch
except ImportError:  # Python 2
    from mock import patch

@patch('splitwise.Splitwise._Splitwise__makeRequest')
class CreateCommentTestCase(unittest.TestCase):

    def setUp(self):
        self.sObj = Splitwise('consumerkey', 'consumersecret')

    def test_createComment_success(self, mockMakeRequest):
        mockMakeRequest.return_value = '{"comment": {"id": 211110000, "content": "Test for create comment", "comment_type": "User", "relation_type": "ExpenseComment", "relation_id": 982430660, "created_at": "2020-08-09T16:14:52Z", "deleted_at": null, "user": {"id": 784241, "first_name": "ruks", "last_name": null, "picture": {"medium": "https://s3.amazonaws.com/splitwise/uploads/user/default_avatars/avatar-ruby47-100px.png"}}}, "errors":{}}'  # noqa: E501
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'
        comment, errors = self.sObj.createComment(data)
        mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                           "POST",
                                           data)
        self.assertEqual(comment.getId(), 211110000)
        self.assertEqual(comment.getContent(), 'Test for create comment')
        self.assertEqual(comment.getCommentType(), 'User')
        self.assertEqual(comment.getRelationType(), 'ExpenseComment')
        self.assertEqual(comment.getRelationId(), 982430660)
        self.assertEqual(comment.getCreatedAt(), '2020-08-09T16:14:52Z')
        self.assertEqual(comment.getDeletedAt(), None)
        user = comment.getCommentedUser()
        self.assertEqual(user.getFirstName(), 'ruks')
        self.assertEqual(user.getLastName(), None)

    def test_createComment_error(self, mockMakeRequest):
        mockMakeRequest.return_value = '{"comment":{},"errors":{"base":["An unknown error occurred. Please try again or contact support@splitwise.com if you experience repeated issues. Sorry for the trouble!"]}}'  # noqa: E501
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'

        comment, errors = self.sObj.createComment(data)
        mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                           "POST",
                                           data)
        self.assertEqual(errors.getErrors(), {
            'base': [
                'An unknown error occurred. Please try again or contact support@splitwise.com if you experience repeated issues. \
Sorry for the trouble!'
            ]})

    def test_createComment_exception(self, mockMakeRequest):
        mockMakeRequest.side_effect = Exception(
            "Invalid response %s. Please check your consumer key and secret." % 404)
        data = '{"expense_id": 982430660, "content": "Test for create comment"}'
        with self.assertRaises(Exception):
            comment, errors = self.sObj.createComment(data)
            mockMakeRequest.assert_called_with("https://secure.splitwise.com/api/v3.0/create_comment",
                                               "POST",
                                               data)
<|endoftext|>


<commit_before>from functools import partial
from enum import Enum

import requests

try:
	import urlparse as parse
except:
	from urllib import parse

class InvalidEndpointException(Exception):
	pass

class InvalidHookException(Exception):
	pass

class Hooks(Enum):
	request_created_hook = 1

	@classmethod
	def has_value(cls, value):
		return any(value == item.value for item in cls)

	@classmethod
	def has_name(cls, name):
		return any(name == item.name for item in cls)

class BaseClient(object):

	endpoints = []
	separator = '/'
	base_url = ''

	def __init__(self, path='', path_params=None, **kwargs):
		self.path = path
		self.path_params = path_params
		if not self.path_params:
			self.path_params = {}

		self.kwargs = kwargs
		self.__define_convenience_methods()

		self._create_request()
		self._run_hook(Hooks.request_created_hook.name)

	def __call__(self, value):
		path_param_key = self.path.split(self.separator)[-1]
		self.path_params[path_param_key] = str(value)
		return self

	def __getattr__(self, name):
		if not Hooks.has_name(name):
			new_path = self.separator.join((self.path, name)) if self.path else name
			return self.__class__(new_path, self.path_params, **self.kwargs)
		return object.__getattribute__(self, name)

	def _run_hook(self, name):
		if hasattr(self, name):
			getattr(self, name)()

	def _find_endpoint(self, method):
		endpoint = None
		to_match = '{} {}'.format(method, self.path)

		for e in self.endpoints:
			if e.match(to_match):
				endpoint = e

		if not endpoint:
			raise InvalidEndpointException('{} is not a valid endpoint.'.format(to_match))

		return endpoint

	def _create_request(self):
		self.request = requests.Request()

	def _validate_endpoint(self, endpoint, params):
		endpoint.validate(params, self.path_params)

	def _build_url(self, path):
		url = parse.urljoin(self.base_url, path)
		for param, value in self.path_params.items():
			url = url.replace('<{}>'.format(param), value)
		return url

	def _finalize_request(self, method, payload):
		endpoint = self._find_endpoint(method)
		self._validate_endpoint(endpoint, payload)

		self.request.url = self._build_url(endpoint.path)
		self.request.data = payload
		self.request.method = method

	def _send_request(self, method, **kwargs):
		self._finalize_request(method, kwargs)
		return requests.session().send(self.request.prepare())

	def __define_convenience_methods(self):
		actions = ['POST', 'GET', 'PUT', 'PATCH', 'DELETE', 'OPTIONS', 'HEAD']

		for action in actions:
			setattr(self, action.lower(), partial(self._send_request, action))

<commit_msg>Remove unneeded code<commit_after>from functools import partial
from enum import Enum

import requests

try:
	import urlparse as parse
except:
	from urllib import parse

class InvalidEndpointException(Exception):
	pass

class Hooks(Enum):
	request_created_hook = 1

	@classmethod
	def has_value(cls, value):
		return any(value == item.value for item in cls)

	@classmethod
	def has_name(cls, name):
		return any(name == item.name for item in cls)

class BaseClient(object):

	endpoints = []
	separator = '/'
	base_url = ''

	def __init__(self, path='', path_params=None, **kwargs):
		self.path = path
		self.path_params = path_params
		if not self.path_params:
			self.path_params = {}

		self.kwargs = kwargs
		self.__define_convenience_methods()

		self._create_request()
		self._run_hook(Hooks.request_created_hook.name)

	def __call__(self, value):
		path_param_key = self.path.split(self.separator)[-1]
		self.path_params[path_param_key] = str(value)
		return self

	def __getattr__(self, name):
		if not Hooks.has_name(name):
			new_path = self.separator.join((self.path, name)) if self.path else name
			return self.__class__(new_path, self.path_params, **self.kwargs)
		return object.__getattribute__(self, name)

	def _run_hook(self, name):
		if hasattr(self, name):
			getattr(self, name)()

	def _find_endpoint(self, method):
		endpoint = None
		to_match = '{} {}'.format(method, self.path)

		for e in self.endpoints:
			if e.match(to_match):
				endpoint = e

		if not endpoint:
			raise InvalidEndpointException('{} is not a valid endpoint.'.format(to_match))

		return endpoint

	def _create_request(self):
		self.request = requests.Request()

	def _validate_endpoint(self, endpoint, params):
		endpoint.validate(params, self.path_params)

	def _build_url(self, path):
		url = parse.urljoin(self.base_url, path)
		for param, value in self.path_params.items():
			url = url.replace('<{}>'.format(param), value)
		return url

	def _finalize_request(self, method, payload):
		endpoint = self._find_endpoint(method)
		self._validate_endpoint(endpoint, payload)

		self.request.url = self._build_url(endpoint.path)
		self.request.data = payload
		self.request.method = method

	def _send_request(self, method, **kwargs):
		self._finalize_request(method, kwargs)
		return requests.session().send(self.request.prepare())

	def __define_convenience_methods(self):
		actions = ['POST', 'GET', 'PUT', 'PATCH', 'DELETE', 'OPTIONS', 'HEAD']

		for action in actions:
			setattr(self, action.lower(), partial(self._send_request, action))

<|endoftext|>


<commit_before>
from django.forms import widgets
from django import models
from django.conf import settings

class RichTextField(models.TextField):
    """
    Field is used just to distinguish a rich text field.
    This way, if you'd like to use the Ckeditor in your
    admin, you can just add the following to your ModelAdmin:

    formfield_overrides = {
        models.RichTextField: {'widget': fields.CkeditorWidget},
    }

    and your richtext field will use the Ckeditor.
    """
    pass

class CkeditorWidget(widgets.Textarea):
    """Ckeditor is a richtext editor. Most basic usage requires a class
    of ckeditor so the javascript knows what to change."""

    def __init__(self, attrs=None):
        default_attrs = {'class':"ckeditor"}
        if attrs:
            default_attrs.update(attrs)
        super(CkeditorWidget, self).__init__(default_attrs)

    class Media:
        css = {'all': (settings.STATIC_URL + "ckeditor/ckeditor.css",)}
        js = (settings.STATIC_URL + "ckeditor/ckeditor.js",)<commit_msg>fixed bad import<commit_after>
from django.forms import widgets
from django.db import models
from django.conf import settings

class RichTextField(models.TextField):
    """
    Field is used just to distinguish a rich text field.
    This way, if you'd like to use the Ckeditor in your
    admin, you can just add the following to your ModelAdmin:

    formfield_overrides = {
        models.RichTextField: {'widget': fields.CkeditorWidget},
    }

    and your richtext field will use the Ckeditor.
    """
    pass

class CkeditorWidget(widgets.Textarea):
    """Ckeditor is a richtext editor. Most basic usage requires a class
    of ckeditor so the javascript knows what to change."""

    def __init__(self, attrs=None):
        default_attrs = {'class':"ckeditor"}
        if attrs:
            default_attrs.update(attrs)
        super(CkeditorWidget, self).__init__(default_attrs)

    class Media:
        css = {'all': (settings.STATIC_URL + "ckeditor/ckeditor.css",)}
        js = (settings.STATIC_URL + "ckeditor/ckeditor.js",)<|endoftext|>


<commit_before>from datetime import datetime

from django.conf import settings
from django.db import models

import amo
from translations.fields import TranslatedField, translations_with_fallback

class Addon(amo.ModelBase):
    STATUS_CHOICES = (
            (0, 'Null'),
            (1, 'In the sandbox'),
            (2, 'Pending approval'),
            (3, 'Nominated to be public'),
            (4, 'Public'),
            (5, 'Disabled'),
            (6, 'Listed'),
            (7, 'Beta'),
            )

    CONTRIBUTIONS_CHOICES = (
            (0, 'None'),
            (1, 'Passive; user shown message next to download button'),
            (2, 'User shown splash screen after download'),
            (3, 'Roadblock; User shown splash screen before download'),
    )

    guid = models.CharField(max_length=255, unique=True)
    name = TranslatedField()
    defaultlocale = models.CharField(max_length=10,
                                     default=settings.LANGUAGE_CODE)

    addontype = models.ForeignKey('AddonType')
    status = models.PositiveIntegerField(choices=STATUS_CHOICES, db_index=True)
    higheststatus = models.PositiveIntegerField(choices=STATUS_CHOICES,
                    help_text="An upper limit for what an author can change.")
    icontype = models.CharField(max_length=25, blank=True)
    homepage = TranslatedField()
    supportemail = TranslatedField()
    supporturl = TranslatedField()
    description = TranslatedField()
    summary = TranslatedField()
    developercomments = TranslatedField()
    eula = TranslatedField()
    privacypolicy = TranslatedField()
    the_reason = TranslatedField()
    the_future = TranslatedField()

    averagerating = models.CharField(max_length=255, default=0)
    bayesianrating = models.FloatField(default=0, db_index=True)
    totalreviews = models.PositiveIntegerField(default=0)
    weeklydownloads = models.PositiveIntegerField(default=0)
    totaldownloads = models.PositiveIntegerField(default=0)
    average_daily_downloads = models.PositiveIntegerField(default=0)
    average_daily_users = models.PositiveIntegerField(default=0)
    sharecount = models.PositiveIntegerField(default=0, db_index=True)

    inactive = models.BooleanField(default=False, db_index=True)
    trusted = models.BooleanField(default=False)
    viewsource = models.BooleanField(default=False)
    publicstats = models.BooleanField(default=False)
    prerelease = models.BooleanField(default=False)
    adminreview = models.BooleanField(default=False)
    sitespecific = models.BooleanField(default=False)
    externalsoftware = models.BooleanField(default=False)
    binary = models.BooleanField(default=False,
                            help_text="Does the add-on contain a binary?")
    dev_agreement = models.BooleanField(default=False,
                            help_text="Has the dev agreement been signed?")
    wants_contributions = models.BooleanField(default=False)
    show_beta = models.BooleanField(default=True)

    nominationdate = models.DateTimeField(null=True)
    target_locale = models.CharField(max_length=255, db_index=True, blank=True,
                            help_text="For dictionaries and language packs")
    locale_disambiguation = models.CharField(max_length=255, blank=True,
                            help_text="For dictionaries and language packs")

    paypal_id = models.CharField(max_length=255, blank=True)
    suggested_amount = models.CharField(max_length=255, blank=True,
                            help_text="Requested donation amount.")
    annoying = models.PositiveIntegerField(choices=STATUS_CHOICES, default=0)

    get_satisfaction_company = models.CharField(max_length=255, blank=True)
    get_satisfaction_product = models.CharField(max_length=255, blank=True)

    users = models.ManyToManyField('users.User')

    class Meta:
        db_table = 'addons'

    def __unicode__(self):
        return '%s: %s' % (self.id, self.name)

    def get_absolute_url(self):
        # XXX: use reverse
        return '/addon/%s' % self.id

    def fetch_translations(self, ids, lang):
        return translations_with_fallback(ids, lang, self.defaultlocale)

class AddonType(amo.ModelBase):

    name = TranslatedField()
    name_plural = TranslatedField()
    description = TranslatedField()

    class Meta:
        db_table = 'addontypes'

    def __unicode__(self):
        return unicode(self.name)

class BlacklistedGuid(models.Model):

    guid = models.CharField(max_length=255, unique=True)
    created = models.DateTimeField(default=datetime.now, editable=False)

    class Meta:
        db_table = 'blacklisted_guids'

    def __unicode__(self):
        return self.guid

class Feature(amo.ModelBase):

    addon = models.ForeignKey(Addon)
    start = models.DateTimeField()
    end = models.DateTimeField()
    locale = models.CharField(max_length=10, default='', blank=True)
    application = models.ForeignKey('amo.Application')

    class Meta:
        db_table = 'features'

    def __unicode__(self):
        return '%s (%s: %s)' % (self.addon.name, self.application.name,
                                self.locale)
<commit_msg>annoying should use contrib choices<commit_after>from datetime import datetime

from django.conf import settings
from django.db import models

import amo
from translations.fields import TranslatedField, translations_with_fallback

class Addon(amo.ModelBase):
    STATUS_CHOICES = (
            (0, 'Null'),
            (1, 'In the sandbox'),
            (2, 'Pending approval'),
            (3, 'Nominated to be public'),
            (4, 'Public'),
            (5, 'Disabled'),
            (6, 'Listed'),
            (7, 'Beta'),
            )

    CONTRIBUTIONS_CHOICES = (
            (0, 'None'),
            (1, 'Passive; user shown message next to download button'),
            (2, 'User shown splash screen after download'),
            (3, 'Roadblock; User shown splash screen before download'),
    )

    guid = models.CharField(max_length=255, unique=True)
    name = TranslatedField()
    defaultlocale = models.CharField(max_length=10,
                                     default=settings.LANGUAGE_CODE)

    addontype = models.ForeignKey('AddonType')
    status = models.PositiveIntegerField(choices=STATUS_CHOICES, db_index=True)
    higheststatus = models.PositiveIntegerField(choices=STATUS_CHOICES,
                    help_text="An upper limit for what an author can change.")
    icontype = models.CharField(max_length=25, blank=True)
    homepage = TranslatedField()
    supportemail = TranslatedField()
    supporturl = TranslatedField()
    description = TranslatedField()
    summary = TranslatedField()
    developercomments = TranslatedField()
    eula = TranslatedField()
    privacypolicy = TranslatedField()
    the_reason = TranslatedField()
    the_future = TranslatedField()

    averagerating = models.CharField(max_length=255, default=0)
    bayesianrating = models.FloatField(default=0, db_index=True)
    totalreviews = models.PositiveIntegerField(default=0)
    weeklydownloads = models.PositiveIntegerField(default=0)
    totaldownloads = models.PositiveIntegerField(default=0)
    average_daily_downloads = models.PositiveIntegerField(default=0)
    average_daily_users = models.PositiveIntegerField(default=0)
    sharecount = models.PositiveIntegerField(default=0, db_index=True)

    inactive = models.BooleanField(default=False, db_index=True)
    trusted = models.BooleanField(default=False)
    viewsource = models.BooleanField(default=False)
    publicstats = models.BooleanField(default=False)
    prerelease = models.BooleanField(default=False)
    adminreview = models.BooleanField(default=False)
    sitespecific = models.BooleanField(default=False)
    externalsoftware = models.BooleanField(default=False)
    binary = models.BooleanField(default=False,
                            help_text="Does the add-on contain a binary?")
    dev_agreement = models.BooleanField(default=False,
                            help_text="Has the dev agreement been signed?")
    wants_contributions = models.BooleanField(default=False)
    show_beta = models.BooleanField(default=True)

    nominationdate = models.DateTimeField(null=True)
    target_locale = models.CharField(max_length=255, db_index=True, blank=True,
                            help_text="For dictionaries and language packs")
    locale_disambiguation = models.CharField(max_length=255, blank=True,
                            help_text="For dictionaries and language packs")

    paypal_id = models.CharField(max_length=255, blank=True)
    suggested_amount = models.CharField(max_length=255, blank=True,
                            help_text="Requested donation amount.")
    annoying = models.PositiveIntegerField(choices=CONTRIBUTIONS_CHOICES,
                                           default=0)

    get_satisfaction_company = models.CharField(max_length=255, blank=True)
    get_satisfaction_product = models.CharField(max_length=255, blank=True)

    users = models.ManyToManyField('users.User')

    class Meta:
        db_table = 'addons'

    def __unicode__(self):
        return '%s: %s' % (self.id, self.name)

    def get_absolute_url(self):
        # XXX: use reverse
        return '/addon/%s' % self.id

    def fetch_translations(self, ids, lang):
        return translations_with_fallback(ids, lang, self.defaultlocale)

class AddonType(amo.ModelBase):

    name = TranslatedField()
    name_plural = TranslatedField()
    description = TranslatedField()

    class Meta:
        db_table = 'addontypes'

    def __unicode__(self):
        return unicode(self.name)

class BlacklistedGuid(models.Model):

    guid = models.CharField(max_length=255, unique=True)
    created = models.DateTimeField(default=datetime.now, editable=False)

    class Meta:
        db_table = 'blacklisted_guids'

    def __unicode__(self):
        return self.guid

class Feature(amo.ModelBase):

    addon = models.ForeignKey(Addon)
    start = models.DateTimeField()
    end = models.DateTimeField()
    locale = models.CharField(max_length=10, default='', blank=True)
    application = models.ForeignKey('amo.Application')

    class Meta:
        db_table = 'features'

    def __unicode__(self):
        return '%s (%s: %s)' % (self.addon.name, self.application.name,
                                self.locale)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################
from openerp.osv import fields, osv

class product_template(osv.osv):
    _inherit = "product.template"
    _columns = {
        'email_template_id': fields.many2one('email.template','Product Email Template'),
    }

class email_template(osv.osv):
    _inherit = 'email.template'

    def default_get(self, cr, uid, fields, context=None):
        res = super(email_template, self).default_get(cr, uid, fields, context)
        if context.get('form_view_ref') == 'account_product_template.view_email_template_form_product':
            res['email_from'] = '${(user.email)|safe}'
            res['email_to'] = '${(object.partner_id.email)|safe}'
            res['partner_to'] = '${object.partner_id.id}'
            res['name'] = context.get('product_name')
            res['subject'] = context.get('product_name')
            res['model_id'] = self.pool.get('ir.model').search(cr, uid, [('model', '=', 'account.invoice')], context=context)
        return res

class account_invoice(osv.Model):
    _inherit = 'account.invoice'

    def invoice_validate(self, cr, uid, ids, context=None):
        if context is None:
            context = {}
        mail_obj = self.pool.get('mail.compose.message')
        res = super(account_invoice, self).invoice_validate(cr, uid, ids, context=context)
        for line in self.browse(cr, uid, ids[0], context=context).invoice_line:
            if line.product_id.email_template_id:
                message_wiz_id = mail_obj.create(cr, uid, {
                    'model': 'account.invoice',
                    'res_id': ids[0],
                    'template_id': line.product_id.email_template_id.id,
                    'body': line.product_id.email_template_id.body_html
                }, context=context)
                mail_obj.send_mail(cr, uid, [message_wiz_id], context=context)
        return res

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<commit_msg>[IMP] add partner as a follower first.<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2010 Tiny SPRL (<http://tiny.be>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################
from openerp.osv import fields, osv

class product_template(osv.osv):
    _inherit = "product.template"
    _columns = {
        'email_template_id': fields.many2one('email.template','Product Email Template'),
    }

class email_template(osv.osv):
    _inherit = 'email.template'

    def default_get(self, cr, uid, fields, context=None):
        res = super(email_template, self).default_get(cr, uid, fields, context)
        if context.get('form_view_ref') == 'account_product_template.view_email_template_form_product':
            res['email_from'] = '${(user.email)|safe}'
            res['email_to'] = '${(object.partner_id.email)|safe}'
            res['partner_to'] = '${object.partner_id.id}'
            res['name'] = context.get('product_name')
            res['subject'] = context.get('product_name')
            res['model_id'] = self.pool.get('ir.model').search(cr, uid, [('model', '=', 'account.invoice')], context=context)
        return res

class account_invoice(osv.Model):
    _inherit = 'account.invoice'

    def invoice_validate(self, cr, uid, ids, context=None):
        if context is None:
            context = {}
        mail_obj = self.pool.get('mail.compose.message')
        res = super(account_invoice, self).invoice_validate(cr, uid, ids, context=context)
        invoice = self.browse(cr, uid, ids[0], context=context)
        for line in invoice.invoice_line:
            # fetch the partner's id and subscribe the partner to the invoice
            if invoice.partner_id.id not in invoice.message_follower_ids:
                self.message_subscribe(cr, uid, [invoice.id], [invoice.partner_id.id], context=context)
            if line.product_id.email_template_id:
                message_wiz_id = mail_obj.create(cr, uid, {
                    'model': 'account.invoice',
                    'res_id': ids[0],
                    'template_id': line.product_id.email_template_id.id,
                    'body': line.product_id.email_template_id.body_html
                }, context=context)
                mail_obj.send_mail(cr, uid, [message_wiz_id], context=context)
        return res

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:
<|endoftext|>


<commit_before>"""A parser for FAT file system images"""

import io
import posixpath
import shutil
from stat import *
import time

from . import *
from ..io import *
from ..util import *

FatHeader = Struct('FatHeader', [
 ('jump', Struct.STR % 3),
 ('oemName', Struct.STR % 8),
 ('bytesPerSector', Struct.INT16),
 ('sectorsPerCluster', Struct.INT8),
 ('reservedSectors', Struct.INT16),
 ('fatCopies', Struct.INT8),
 ('rootEntries', Struct.INT16),
 ('...', 2),
 ('mediaDescriptor', Struct.INT8),
 ('sectorsPerFat', Struct.INT16),
 ('...', 8),
 ('sectors', Struct.INT32),
 ('...', 2),
 ('extendedSignature', Struct.STR % 1),
 ('serialNumber', Struct.INT32),
 ('volumeLabel', Struct.STR % 11),
 ('fsType', Struct.STR % 8),
 ('...', 448),
 ('signature', Struct.STR % 2),
])
fatHeaderSignature = b'\x55\xaa'
fatHeaderExtendedSignature = b'\x29'

FatDirEntry = Struct('FatDirEntry', [
 ('name', Struct.STR % 8),
 ('ext', Struct.STR % 3),
 ('attr', Struct.INT8),
 ('...', 1),
 ('ctimeCs', Struct.INT8),
 ('...', 8),
 ('time', Struct.INT16),
 ('date', Struct.INT16),
 ('cluster', Struct.INT16),
 ('size', Struct.INT32),
])

VfatDirEntry = Struct('VfatDirEntry', [
 ('sequence', Struct.INT8),
 ('name1', Struct.STR % 10),
 ('attr', Struct.INT8),
 ('...', 1),
 ('checksum', Struct.INT8),
 ('name2', Struct.STR % 12),
 ('...', 2),
 ('name3', Struct.STR % 4),
])

def isFat(file):
 header = FatHeader.unpack(file)
 return header and header.signature == fatHeaderSignature and header.extendedSignature == fatHeaderExtendedSignature and header.fsType.startswith(b'FAT')

def readFat(file):
 header = FatHeader.unpack(file)

 if header.signature != fatHeaderSignature or header.extendedSignature != fatHeaderExtendedSignature:
  raise Exception('Wrong magic')

 fatOffset = header.reservedSectors * header.bytesPerSector
 rootOffset = fatOffset + header.fatCopies * header.sectorsPerFat * header.bytesPerSector
 dataOffset = rootOffset + ((header.rootEntries * FatDirEntry.size - 1) // header.bytesPerSector + 1) * header.bytesPerSector

 file.seek(fatOffset)
 if header.fsType == b'FAT12   ':
  endMarker = 0xfff
  packedClusters = [parse32le(file.read(3) + b'\0') for i in range(0, header.sectorsPerFat * header.bytesPerSector, 3)]
  clusters = [cluster for packed in packedClusters for cluster in [packed & 0xfff, (packed >> 12) & 0xfff]]
 elif header.fsType == b'FAT16   ':
  endMarker = 0xffff
  clusters = [parse16le(file.read(2)) for i in range(0, header.sectorsPerFat * header.bytesPerSector, 2)]
 else:
  raise Exception('Unknown FAT width')

 def readDir(entries, path=''):
  offset = 0
  vfatName = b''
  while entries[offset:offset+1] != b'\0':
   entry = FatDirEntry.unpack(entries, offset)
   if entry.name[0:1] != b'\xe5':
    if entry.attr == 0x0f:
     # VFAT
     vfatEntry = VfatDirEntry.unpack(entries, offset)
     vfatName = vfatEntry.name1 + vfatEntry.name2 + vfatEntry.name3 + vfatName
    else:
     if vfatName != b'':
      name = vfatName.decode('utf16').rstrip(u'\0\uffff')
      vfatName = b''
     else:
      name = entry.name.decode('ascii').rstrip(' ')
      if name[0] == '\x05':
       name = '\xe5' + name[1:]
      ext = entry.ext.decode('ascii').rstrip(' ')
      if ext != '':
       name += '.' + ext

     if name != '.' and name != '..':
      isLink = (entry.attr & 0x04) and (entry.ctimeCs & 0xe1) == 0x21
      isDir = entry.attr & 0x10

      def generateChunks(cluster=entry.cluster, size=entry.size, isDir=isDir):
       read = 0
       while cluster != 0 and cluster != endMarker and (read < size or isDir):
        file.seek(dataOffset + (cluster - 2) * header.sectorsPerCluster * header.bytesPerSector)
        block = file.read(header.sectorsPerCluster * header.bytesPerSector)
        yield block if isDir else block[:size-read]
        read += len(block)
        cluster = clusters[cluster]

      contents = ChunkedFile(generateChunks, entry.size if not isDir else -1)
      yield UnixFile(
       path = path + '/' + name,
       size = entry.size,
       mtime = time.mktime((1980 + (entry.date >> 9), (entry.date >> 5) & 0xf, entry.date & 0x1f, entry.time >> 11, (entry.time >> 5) & 0x3f, (entry.time & 0x1f) * 2, -1, -1, -1)),
       mode = S_IFDIR if isDir else S_IFLNK if isLink else S_IFREG,
       uid = 0,
       gid = 0,
       contents = contents if not isDir else None,
      )

      if isDir:
       for f in readDir(contents.read(), path + '/' + name):
        yield f

   offset += FatDirEntry.size

 file.seek(rootOffset)
 for f in readDir(file.read(dataOffset - rootOffset)):
  yield f

def writeFat(files, size, outFile):
 files = {f.path: f for f in files}
 tree = {'': set()}
 for path in files:
  while path != '':
   parent = posixpath.dirname(path).rstrip('/')
   tree.setdefault(parent, set()).add(path)
   path = parent

 sectorSize = 0x200
 clusterSize = 0x4000

 sectors = size // sectorSize
 fatSize = (size // clusterSize + 1) // 2 * 3
 fatSectors = (fatSize + sectorSize - 1) // sectorSize

 outFile.write(FatHeader.pack(
  jump = b'\xeb\0\x90',
  oemName = 8*b'\0',
  bytesPerSector = sectorSize,
  sectorsPerCluster = clusterSize // sectorSize,
  reservedSectors = 1,
  fatCopies = 1,
  rootEntries = clusterSize // FatDirEntry.size,
  sectors = sectors,
  mediaDescriptor = 0xf8,
  sectorsPerFat = fatSectors,
  extendedSignature = fatHeaderExtendedSignature,
  serialNumber = 0,
  volumeLabel = 11*b' ',
  fsType = b'FAT12   ',
  signature = fatHeaderSignature,
 ))
 for i in range(sectors - 1):
  outFile.write(sectorSize * b'\0')

 fatOffset = sectorSize
 rootOffset = fatOffset + fatSectors * sectorSize
 dataOffset = rootOffset + clusterSize

 clusters = [0xff8, 0xfff]
 def writeData(f):
  f.seek(0)
  outFile.seek(dataOffset + (len(clusters) - 2) * clusterSize)
  shutil.copyfileobj(f, outFile)
  nc = (f.tell() + clusterSize - 1) // clusterSize
  for i in range(nc):
   clusters.append(len(clusters) + 1 if i < nc-1 else 0xfff)
  return (len(clusters)-nc if nc else 0), f.tell()

 def dirEntries(pc, c):
  return FatDirEntry.pack(
    name = b'.       ',
    ext = b'   ',
    attr = 0x10,
    ctimeCs = 0,
    time = 0,
    date = 0,
    cluster = c,
    size = 0,
   ) + FatDirEntry.pack(
    name = b'..      ',
    ext = b'   ',
    attr = 0x10,
    ctimeCs = 0,
    time = 0,
    date = 0,
    cluster = pc,
    size = 0,
   )

 dirs = {}
 def writeDir(path):
  data = io.BytesIO()
  if path != '':
   data.write(dirEntries(0, 0))
  for p in tree.get(path, set()):
   file = files.get(p, UnixFile(p, 0, 0, S_IFDIR | 0o775, 0, 0, None))
   c, s = writeData(file.contents if not S_ISDIR(file.mode) else writeDir(file.path))
   if S_ISDIR(file.mode):
    dirs[file.path] = c

   name, ext = (posixpath.basename(file.path).upper() + '.').split('.', 1)
   name = name[:8].ljust(8, ' ').encode('ascii')
   ext = ext[:3].ljust(3, ' ').encode('ascii')
   sum = 0
   for chr in (name + ext):
    sum = (((sum & 1) << 7) + (sum >> 1) + chr) & 0xff

   fn = posixpath.basename(file.path) + '\0'
   vfatEntries = [fn[o:o+13] for o in range(0, len(fn), 13)]
   for i, n in list(enumerate(vfatEntries))[::-1]:
    n = n.encode('utf-16le').ljust(26, b'\xff')
    data.write(VfatDirEntry.pack(
     sequence = i + 1 + (0x40 if i == len(vfatEntries)-1 else 0),
     name1 = n[:10],
     attr = 0x0f,
     checksum = sum,
     name2 = n[10:22],
     name3 = n[22:],
    ))

   t = time.localtime(file.mtime)
   data.write(FatDirEntry.pack(
    name = name,
    ext = ext,
    attr = 0x10 if S_ISDIR(file.mode) else 0x04 if S_ISLNK(file.mode) else 0,
    ctimeCs = 0x21 if S_ISLNK(file.mode) else 0,
    time = (t.tm_hour << 11) + (t.tm_min << 5) + t.tm_sec // 2,
    date = (max(t.tm_year - 1980, 0) << 9) + (t.tm_mon << 5) + t.tm_mday,
    cluster = c,
    size = s if not S_ISDIR(file.mode) else 0,
   ))
  return data

 root = writeDir('')
 root.seek(0)
 outFile.seek(rootOffset)
 shutil.copyfileobj(root, outFile)

 for p, c in dirs.items():
  parent = posixpath.split(p)[0]
  outFile.seek(dataOffset + (c - 2) * clusterSize)
  outFile.write(dirEntries(dirs[parent] if parent != '/' else 0, c))

 outFile.seek(fatOffset)
 for i in range(0, len(clusters), 2):
  outFile.write(dump32le(clusters[i] + ((clusters[i+1] << 12) if i+1 < len(clusters) else 0))[:3])
<commit_msg>FAT: Fix deleted file handling<commit_after>"""A parser for FAT file system images"""

import io
import posixpath
import shutil
from stat import *
import time

from . import *
from ..io import *
from ..util import *

FatHeader = Struct('FatHeader', [
 ('jump', Struct.STR % 3),
 ('oemName', Struct.STR % 8),
 ('bytesPerSector', Struct.INT16),
 ('sectorsPerCluster', Struct.INT8),
 ('reservedSectors', Struct.INT16),
 ('fatCopies', Struct.INT8),
 ('rootEntries', Struct.INT16),
 ('...', 2),
 ('mediaDescriptor', Struct.INT8),
 ('sectorsPerFat', Struct.INT16),
 ('...', 8),
 ('sectors', Struct.INT32),
 ('...', 2),
 ('extendedSignature', Struct.STR % 1),
 ('serialNumber', Struct.INT32),
 ('volumeLabel', Struct.STR % 11),
 ('fsType', Struct.STR % 8),
 ('...', 448),
 ('signature', Struct.STR % 2),
])
fatHeaderSignature = b'\x55\xaa'
fatHeaderExtendedSignature = b'\x29'

FatDirEntry = Struct('FatDirEntry', [
 ('name', Struct.STR % 8),
 ('ext', Struct.STR % 3),
 ('attr', Struct.INT8),
 ('...', 1),
 ('ctimeCs', Struct.INT8),
 ('...', 8),
 ('time', Struct.INT16),
 ('date', Struct.INT16),
 ('cluster', Struct.INT16),
 ('size', Struct.INT32),
])

VfatDirEntry = Struct('VfatDirEntry', [
 ('sequence', Struct.INT8),
 ('name1', Struct.STR % 10),
 ('attr', Struct.INT8),
 ('...', 1),
 ('checksum', Struct.INT8),
 ('name2', Struct.STR % 12),
 ('...', 2),
 ('name3', Struct.STR % 4),
])

def isFat(file):
 header = FatHeader.unpack(file)
 return header and header.signature == fatHeaderSignature and header.extendedSignature == fatHeaderExtendedSignature and header.fsType.startswith(b'FAT')

def readFat(file):
 header = FatHeader.unpack(file)

 if header.signature != fatHeaderSignature or header.extendedSignature != fatHeaderExtendedSignature:
  raise Exception('Wrong magic')

 fatOffset = header.reservedSectors * header.bytesPerSector
 rootOffset = fatOffset + header.fatCopies * header.sectorsPerFat * header.bytesPerSector
 dataOffset = rootOffset + ((header.rootEntries * FatDirEntry.size - 1) // header.bytesPerSector + 1) * header.bytesPerSector

 file.seek(fatOffset)
 if header.fsType == b'FAT12   ':
  endMarker = 0xfff
  packedClusters = [parse32le(file.read(3) + b'\0') for i in range(0, header.sectorsPerFat * header.bytesPerSector, 3)]
  clusters = [cluster for packed in packedClusters for cluster in [packed & 0xfff, (packed >> 12) & 0xfff]]
 elif header.fsType == b'FAT16   ':
  endMarker = 0xffff
  clusters = [parse16le(file.read(2)) for i in range(0, header.sectorsPerFat * header.bytesPerSector, 2)]
 else:
  raise Exception('Unknown FAT width')

 def readDir(entries, path=''):
  offset = 0
  vfatName = b''
  while entries[offset:offset+1] != b'\0':
   entry = FatDirEntry.unpack(entries, offset)
   if entry.name[0:1] == b'\xe5':
    vfatName = b''
   else:
    if entry.attr == 0x0f:
     # VFAT
     vfatEntry = VfatDirEntry.unpack(entries, offset)
     vfatName = vfatEntry.name1 + vfatEntry.name2 + vfatEntry.name3 + vfatName
    else:
     if vfatName != b'':
      name = vfatName.decode('utf16').rstrip(u'\0\uffff')
      vfatName = b''
     else:
      name = entry.name.decode('ascii').rstrip(' ')
      if name[0] == '\x05':
       name = '\xe5' + name[1:]
      ext = entry.ext.decode('ascii').rstrip(' ')
      if ext != '':
       name += '.' + ext

     if name != '.' and name != '..':
      isLink = (entry.attr & 0x04) and (entry.ctimeCs & 0xe1) == 0x21
      isDir = entry.attr & 0x10

      def generateChunks(cluster=entry.cluster, size=entry.size, isDir=isDir):
       read = 0
       while cluster != 0 and cluster != endMarker and (read < size or isDir):
        file.seek(dataOffset + (cluster - 2) * header.sectorsPerCluster * header.bytesPerSector)
        block = file.read(header.sectorsPerCluster * header.bytesPerSector)
        yield block if isDir else block[:size-read]
        read += len(block)
        cluster = clusters[cluster]

      contents = ChunkedFile(generateChunks, entry.size if not isDir else -1)
      yield UnixFile(
       path = path + '/' + name,
       size = entry.size,
       mtime = time.mktime((1980 + (entry.date >> 9), (entry.date >> 5) & 0xf, entry.date & 0x1f, entry.time >> 11, (entry.time >> 5) & 0x3f, (entry.time & 0x1f) * 2, -1, -1, -1)),
       mode = S_IFDIR if isDir else S_IFLNK if isLink else S_IFREG,
       uid = 0,
       gid = 0,
       contents = contents if not isDir else None,
      )

      if isDir:
       for f in readDir(contents.read(), path + '/' + name):
        yield f

   offset += FatDirEntry.size

 file.seek(rootOffset)
 for f in readDir(file.read(dataOffset - rootOffset)):
  yield f

def writeFat(files, size, outFile):
 files = {f.path: f for f in files}
 tree = {'': set()}
 for path in files:
  while path != '':
   parent = posixpath.dirname(path).rstrip('/')
   tree.setdefault(parent, set()).add(path)
   path = parent

 sectorSize = 0x200
 clusterSize = 0x4000

 sectors = size // sectorSize
 fatSize = (size // clusterSize + 1) // 2 * 3
 fatSectors = (fatSize + sectorSize - 1) // sectorSize

 outFile.write(FatHeader.pack(
  jump = b'\xeb\0\x90',
  oemName = 8*b'\0',
  bytesPerSector = sectorSize,
  sectorsPerCluster = clusterSize // sectorSize,
  reservedSectors = 1,
  fatCopies = 1,
  rootEntries = clusterSize // FatDirEntry.size,
  sectors = sectors,
  mediaDescriptor = 0xf8,
  sectorsPerFat = fatSectors,
  extendedSignature = fatHeaderExtendedSignature,
  serialNumber = 0,
  volumeLabel = 11*b' ',
  fsType = b'FAT12   ',
  signature = fatHeaderSignature,
 ))
 for i in range(sectors - 1):
  outFile.write(sectorSize * b'\0')

 fatOffset = sectorSize
 rootOffset = fatOffset + fatSectors * sectorSize
 dataOffset = rootOffset + clusterSize

 clusters = [0xff8, 0xfff]
 def writeData(f):
  f.seek(0)
  outFile.seek(dataOffset + (len(clusters) - 2) * clusterSize)
  shutil.copyfileobj(f, outFile)
  nc = (f.tell() + clusterSize - 1) // clusterSize
  for i in range(nc):
   clusters.append(len(clusters) + 1 if i < nc-1 else 0xfff)
  return (len(clusters)-nc if nc else 0), f.tell()

 def dirEntries(pc, c):
  return FatDirEntry.pack(
    name = b'.       ',
    ext = b'   ',
    attr = 0x10,
    ctimeCs = 0,
    time = 0,
    date = 0,
    cluster = c,
    size = 0,
   ) + FatDirEntry.pack(
    name = b'..      ',
    ext = b'   ',
    attr = 0x10,
    ctimeCs = 0,
    time = 0,
    date = 0,
    cluster = pc,
    size = 0,
   )

 dirs = {}
 def writeDir(path):
  data = io.BytesIO()
  if path != '':
   data.write(dirEntries(0, 0))
  for p in tree.get(path, set()):
   file = files.get(p, UnixFile(p, 0, 0, S_IFDIR | 0o775, 0, 0, None))
   c, s = writeData(file.contents if not S_ISDIR(file.mode) else writeDir(file.path))
   if S_ISDIR(file.mode):
    dirs[file.path] = c

   name, ext = (posixpath.basename(file.path).upper() + '.').split('.', 1)
   name = name[:8].ljust(8, ' ').encode('ascii')
   ext = ext[:3].ljust(3, ' ').encode('ascii')
   sum = 0
   for chr in (name + ext):
    sum = (((sum & 1) << 7) + (sum >> 1) + chr) & 0xff

   fn = posixpath.basename(file.path) + '\0'
   vfatEntries = [fn[o:o+13] for o in range(0, len(fn), 13)]
   for i, n in list(enumerate(vfatEntries))[::-1]:
    n = n.encode('utf-16le').ljust(26, b'\xff')
    data.write(VfatDirEntry.pack(
     sequence = i + 1 + (0x40 if i == len(vfatEntries)-1 else 0),
     name1 = n[:10],
     attr = 0x0f,
     checksum = sum,
     name2 = n[10:22],
     name3 = n[22:],
    ))

   t = time.localtime(file.mtime)
   data.write(FatDirEntry.pack(
    name = name,
    ext = ext,
    attr = 0x10 if S_ISDIR(file.mode) else 0x04 if S_ISLNK(file.mode) else 0,
    ctimeCs = 0x21 if S_ISLNK(file.mode) else 0,
    time = (t.tm_hour << 11) + (t.tm_min << 5) + t.tm_sec // 2,
    date = (max(t.tm_year - 1980, 0) << 9) + (t.tm_mon << 5) + t.tm_mday,
    cluster = c,
    size = s if not S_ISDIR(file.mode) else 0,
   ))
  return data

 root = writeDir('')
 root.seek(0)
 outFile.seek(rootOffset)
 shutil.copyfileobj(root, outFile)

 for p, c in dirs.items():
  parent = posixpath.split(p)[0]
  outFile.seek(dataOffset + (c - 2) * clusterSize)
  outFile.write(dirEntries(dirs[parent] if parent != '/' else 0, c))

 outFile.seek(fatOffset)
 for i in range(0, len(clusters), 2):
  outFile.write(dump32le(clusters[i] + ((clusters[i+1] << 12) if i+1 < len(clusters) else 0))[:3])
<|endoftext|>