<commit_before>#!/usr/bin/python
# File:        transform.py
# Description: transform
# Created:     2017-08-19

import numpy as np, quaternion

from .level import WedgeGS

SIMPLE_SIZE = 64

WEDGE_DEF_ROT = np.quaternion(np.cos(np.pi/4), 0, np.sin(np.pi/4), 0)

def convquat(quat):
    return np.array([quat.x, quat.y, quat.z, quat.w])

def length(vec):
    import numpy as np
    return np.linalg.norm(vec)

def normalized(vec):
    import numpy as np
    return vec / np.linalg.norm(vec)

def rotpoint(rot, point):
    import numpy as np
    return (rot * np.quaternion(0, *point) / rot).imag

def rotpointrev(rot, point):
    import numpy as np
    return (rot.conj() * np.quaternion(0, *point) * rot).imag

def vec_angle(va, vb):
    import numpy as np
    ua = normalized(va)
    ub = normalized(vb)
    return np.arccos(np.clip(np.dot(ua, ub), -1.0, 1.0))

def rtri_to_vers(verts):

    """Calculates the versor that aligns the legs of the triangle
    [(0, 0, 0), (0, 0, -1), (0, 1, 0)] with the given right-angled triangle's legs.

    `verts` is an array of 3-dimensional vertices of length 3.
    The first entry is the vertex of the triangle's right angle."""

    import numpy as np, quaternion
    from numpy import pi, sin, cos, arctan2

    pr, pa, pb = verts
    rot = np.quaternion(1, 0, 0, 0)

    # rotate around y for pa
    var = pa - pr
    ay = arctan2(-var[0], -var[2])
    rot *= np.quaternion(cos(ay/2), 0, sin(ay/2), 0)

    # rotate around x for pa
    vaxr = rotpointrev(rot, var)
    ax = arctan2(vaxr[1], -vaxr[2])
    rot *= np.quaternion(cos(ax/2), sin(ax/2), 0, 0)

    # rotate around z for pb
    vbxr = rotpointrev(rot, pb - pr)
    az = arctan2(-vbxr[0], vbxr[1])
    rot *= np.quaternion(cos(az/2), 0, 0, sin(az/2))

    return rot

def rtri_to_transform(verts, srot=None):

    """Converts the given vetices representing a right triangle to a
    transform for a WedgeGS."""

    import numpy as np, quaternion
    from numpy import pi, sin, cos

    pr, pa, pb = verts

    rot = rtri_to_vers(verts)

    pos = (pa + pb) / 2
    scale = [1e-5, length(pr - pb) / SIMPLE_SIZE, length(pr - pa) / SIMPLE_SIZE]
    return pos, convquat(rot), scale

def create_two_wedges(pmax, pnext, plast, objs, simple_args={}):

    """Creates two wedges for an arbitrary triangle.

    pmax - the vertex with the greatest angle
    pnext, plast - the remaining vertices
    objs - the list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np
    from numpy import dot

    vnm = pmax - pnext
    vnl = plast - pnext

    vr = dot(vnm, vnl) / length(vnl) * normalized(vnl)
    pr = vr + pnext

    transform = rtri_to_transform(np.array([pr, pmax, pnext]))
    objs.append(WedgeGS(transform=transform, **simple_args))

    transform = rtri_to_transform(np.array([pr, plast, pmax]))
    objs.append(WedgeGS(transform=transform, **simple_args))

def create_single_wedge(pr, pa, pb, objs, simple_args={}):

    """Creates a single wedge for the given right triangle.

    pmax - the vertex with the right angle
    pnext, plast - the remaining vertices
    objs - the list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np

    transform = rtri_to_transform(np.array([pr, pa, pb]))
    objs.append(WedgeGS(transform=transform, **simple_args))

def create_triangle_simples(verts, objs, simple_args={}):

    """Creates simples for the given triangle.

    verts - sequence containing the three vertices
    objs - list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np
    from numpy import pi

    pa, pb, pc = verts

    ac = abs(vec_angle(pa - pc, pb - pc))
    ab = abs(vec_angle(pa - pb, pc - pb))
    aa = pi - ac - ab

    imax, amax = max(enumerate([aa, ab, ac]), key=lambda e: e[1])

    pmax = verts[imax]
    pnext = verts[(imax + 1) % 3]
    plast = verts[(imax + 2) % 3]

    if -0.001 < abs(amax) - pi/2 < 0.001:
        # very close to right triangle
        create_single_wedge(pmax, pnext, plast, objs, simple_args=simple_args)
    else:
        # any other triangle
        create_two_wedges(pmax, pnext, plast, objs, simple_args=simple_args)

    # objs.append(WedgeGS(
    #     type='SphereGS',
    #     transform=[pr, (), [.7/SIMPLE_SIZE]*3]))

# vim:set sw=4 ts=8 sts=4 et sr ft=python fdm=marker tw=0:
<commit_msg>transform: reduce copies for rotation<commit_after>#!/usr/bin/python
# File:        transform.py
# Description: transform
# Created:     2017-08-19

import numpy as np, quaternion

from .level import WedgeGS

SIMPLE_SIZE = 64

WEDGE_DEF_ROT = np.quaternion(np.cos(np.pi/4), 0, np.sin(np.pi/4), 0)

def convquat(quat):
    return np.array([quat.x, quat.y, quat.z, quat.w])

def length(vec):
    import numpy as np
    return np.linalg.norm(vec)

def normalized(vec):
    import numpy as np
    return vec / np.linalg.norm(vec)

def rotpoint(rot, point):
    import numpy as np
    res = rot * np.quaternion(0, *point)
    res /= rot
    return res.imag

def rotpointrev(rot, point):
    import numpy as np
    res = rot.conj()
    res *= np.quaternion(0, *point)
    res *= rot
    return res.imag

def vec_angle(va, vb):
    import numpy as np
    ua = normalized(va)
    ub = normalized(vb)
    return np.arccos(np.clip(np.dot(ua, ub), -1.0, 1.0))

def rtri_to_vers(verts):

    """Calculates the versor that aligns the legs of the triangle
    [(0, 0, 0), (0, 0, -1), (0, 1, 0)] with the given right-angled triangle's legs.

    `verts` is an array of 3-dimensional vertices of length 3.
    The first entry is the vertex of the triangle's right angle."""

    import numpy as np, quaternion
    from numpy import pi, sin, cos, arctan2

    pr, pa, pb = verts
    rot = np.quaternion(1, 0, 0, 0)

    # rotate around y for pa
    var = pa - pr
    ay = arctan2(-var[0], -var[2])
    rot *= np.quaternion(cos(ay/2), 0, sin(ay/2), 0)

    # rotate around x for pa
    vaxr = rotpointrev(rot, var)
    ax = arctan2(vaxr[1], -vaxr[2])
    rot *= np.quaternion(cos(ax/2), sin(ax/2), 0, 0)

    # rotate around z for pb
    vbxr = rotpointrev(rot, pb - pr)
    az = arctan2(-vbxr[0], vbxr[1])
    rot *= np.quaternion(cos(az/2), 0, 0, sin(az/2))

    return rot

def rtri_to_transform(verts, srot=None):

    """Converts the given vetices representing a right triangle to a
    transform for a WedgeGS."""

    import numpy as np, quaternion
    from numpy import pi, sin, cos

    pr, pa, pb = verts

    rot = rtri_to_vers(verts)

    pos = (pa + pb) / 2
    scale = [1e-5, length(pr - pb) / SIMPLE_SIZE, length(pr - pa) / SIMPLE_SIZE]
    return pos, convquat(rot), scale

def create_two_wedges(pmax, pnext, plast, objs, simple_args={}):

    """Creates two wedges for an arbitrary triangle.

    pmax - the vertex with the greatest angle
    pnext, plast - the remaining vertices
    objs - the list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np
    from numpy import dot

    vnm = pmax - pnext
    vnl = plast - pnext

    vr = dot(vnm, vnl) / length(vnl) * normalized(vnl)
    pr = vr + pnext

    transform = rtri_to_transform(np.array([pr, pmax, pnext]))
    objs.append(WedgeGS(transform=transform, **simple_args))

    transform = rtri_to_transform(np.array([pr, plast, pmax]))
    objs.append(WedgeGS(transform=transform, **simple_args))

def create_single_wedge(pr, pa, pb, objs, simple_args={}):

    """Creates a single wedge for the given right triangle.

    pmax - the vertex with the right angle
    pnext, plast - the remaining vertices
    objs - the list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np

    transform = rtri_to_transform(np.array([pr, pa, pb]))
    objs.append(WedgeGS(transform=transform, **simple_args))

def create_triangle_simples(verts, objs, simple_args={}):

    """Creates simples for the given triangle.

    verts - sequence containing the three vertices
    objs - list to put the objects into
    simple_args - args to pass to WedgeGS"""

    import numpy as np
    from numpy import pi

    pa, pb, pc = verts

    ac = abs(vec_angle(pa - pc, pb - pc))
    ab = abs(vec_angle(pa - pb, pc - pb))
    aa = pi - ac - ab

    imax, amax = max(enumerate([aa, ab, ac]), key=lambda e: e[1])

    pmax = verts[imax]
    pnext = verts[(imax + 1) % 3]
    plast = verts[(imax + 2) % 3]

    if -0.001 < abs(amax) - pi/2 < 0.001:
        # very close to right triangle
        create_single_wedge(pmax, pnext, plast, objs, simple_args=simple_args)
    else:
        # any other triangle
        create_two_wedges(pmax, pnext, plast, objs, simple_args=simple_args)

    # objs.append(WedgeGS(
    #     type='SphereGS',
    #     transform=[pr, (), [.7/SIMPLE_SIZE]*3]))

# vim:set sw=4 ts=8 sts=4 et sr ft=python fdm=marker tw=0:
<|endoftext|>


<commit_before>
from BasicSpiderConfig import ExRule

class Config:

    list_css_rules = {
        '#live-list-contentbox li': {
            'url': 'a::attr(href)',
            'room_name': 'a::attr(title)',
            'tag': 'span.tag.ellipsis::text',
            'people_count': '.dy-num.fr::text'
        }
    }

    ex_rule = ExRule('http://www.douyu.com/all$', list_css_rules=list_css_rules)

    name='douyu.com'
    allowed_domains=['douyu.com']
    start_urls=['http://www.douyu.com/all']
    ex_rules = []
    follow = False

<commit_msg>add author name<commit_after>
from BasicSpiderConfig import ExRule

class Config:

    list_css_rules = {
        '#live-list-contentbox li': {
            'url': 'a::attr(href)',
            'room_name': 'a::attr(title)',
            'tag': 'span.tag.ellipsis::text',
            'people_count': '.dy-num.fr::text',
            'author': '.dy-name::text'
        }
    }

    ex_rule = ExRule('http://www.douyu.com/all$', list_css_rules=list_css_rules)

    name='douyu.com'
    allowed_domains=['douyu.com']
    start_urls=['http://www.douyu.com/all']
    ex_rules = []
    follow = False

<|endoftext|>


<commit_before>"""j2cl_generate_jsunit_suite build macro

Takes Java source that contains JUnit tests and translates it into a goog.testing.testSuite.

Example using a jsunit_test:

j2cl_library(
    name = "mytests",
    ...
)

j2cl_generate_jsunit_suite(
    name = "CodeGen",
    deps = [":mytests"],
    test_class = "mypackage.MyTest",
)

jsunit_test(
  name = "test",
  srcs = [":CodeGen.js.zip"],
  deps = [":CodeGen"],
  ...
)

"""

load("//build_def:j2cl_library.bzl", "j2cl_library")
load("//build_def:j2cl_util.bzl", "get_java_package")

_TEMPLATE = """
// THIS IS GENERATED CODE. DO NOT EDIT.
// GENERATED FROM //%s/BUILD (target: %s)

package %s;

import com.google.j2cl.junit.apt.J2clTestInput;

/**
 * J2CL java tests.
 *
 * @author generated by j2cl_generate_jsunit_suite.bzl.
 */
@J2clTestInput(%s.class)
public class %s {}
"""

def _generate_test_input(name, test_class):
  java_package = get_java_package(PACKAGE_NAME) or "_default_"
  java_class = name.replace('-', '_').title() + "__generated_j2cl_test_input"
  java_code = _TEMPLATE % (PACKAGE_NAME, name, java_package, test_class, java_class)

  native.genrule(
      name = java_class,
      outs = [java_class + ".java"],
      cmd="echo \"%s\" > $@" % java_code,
      tags=["manual", "notap"],
      testonly = 1,
      local=True,
  )

  return java_class + ".java"

def j2cl_generate_jsunit_suite(name, test_class, deps, tags = []):
  """Macro for cross compiling a JUnit Suite to JavaScript testSuite"""

  test_input = _generate_test_input(name, test_class)
  # This target triggers our Java annotation processor and generates the
  # plumbing between a jsunit_test and transpiled JUnit tests
  # It's outputs are:
  #  - test_summary.json that lists all jsunit test suites
  #  - A .testsuite file for every JUnit test containing a jsunit tests suite
  #  - A Java class (JsType) that contains the bridge between the jsunit test
  #        suite and the JUnit test. This is being used from the .testsuite
  #        file.
  #
  # We separated this out into a separate target so we can have dependencies
  # here that users might not have in their tests (e.g. jsinterop annotations).
  # We need the extra dep here on user provided dependencies since our generated
  # code refers to user written code (test cases).
  # Note that test suites are generated with .testsuite extension to avoid
  # j2cl_library automically including them as source.
  j2cl_library(
      name = name + "_lib",
      srcs = [test_input],
      deps = deps +  [
          "//internal_do_not_use:internal_junit_annotations",
          "//third_party/java/gwt:gwt-jsinterop-annotations-j2cl",
      ],
      _js_deps = ["//javascript/closure/testing:testcase"],
      plugins = ["//internal_do_not_use:junit_processor_plugin"],
      testonly = 1,
      tags = tags,
      generate_build_test=False,
  )

  # The Java annotation processor on the above target generates jsunit suites
  # (.testsuite files), but the same jar file also contains unrelated things
  # (e.g. class files), this genrules takes the jar file as input and creates
  # a new zip file that only contains the generated javascript (.testsuite
  # renamed to .js) and the test_summary.json file.
  # This is the format that jsunit_test will later expect.
  # TODO(goktug): use j2cl_library directly from jsunit_test instead of
  # extracting files from jar (output js zip can include all the required
  # files.)
  out_jar = ":lib" + name + "_lib_java_library.jar"
  native.genrule(
      name=name + "_transpile_gen",
      outs=[name + ".js.zip"],
      cmd="\n".join([
          "unzip -q $(location %s) *.testsuite *.json -d zip_out/" % out_jar,
          "cd zip_out/",
          "for f in $$(find -name *.testsuite); do mv $$f $${f/.testsuite/.js}; done",
          "zip -q -r ../$@ .",
      ]),
      testonly=1,
      tags=["manual", "notap"],
      tools=[out_jar],
      local=True,
  )
<commit_msg>Breaks test build of a depending project<commit_after>"""j2cl_generate_jsunit_suite build macro

Takes Java source that contains JUnit tests and translates it into a goog.testing.testSuite.

Example using a jsunit_test:

j2cl_library(
    name = "mytests",
    ...
)

j2cl_generate_jsunit_suite(
    name = "CodeGen",
    deps = [":mytests"],
    test_class = "mypackage.MyTest",
)

jsunit_test(
  name = "test",
  srcs = [":CodeGen.js.zip"],
  deps = [":CodeGen"],
  ...
)

"""

load("//build_def:j2cl_library.bzl", "j2cl_library")
load("//build_def:j2cl_util.bzl", "get_java_package")

_TEMPLATE = """
// THIS IS GENERATED CODE. DO NOT EDIT.
// GENERATED FROM //%s/BUILD (target: %s)

package %s;

import com.google.j2cl.junit.apt.J2clTestInput;

/**
 * J2CL java tests.
 *
 * @author generated by j2cl_generate_jsunit_suite.bzl.
 */
@J2clTestInput(%s.class)
public class %s {}
"""

def _generate_test_input(name, test_class):
  java_package = get_java_package(PACKAGE_NAME) or "_default_"
  java_class = name.replace('-', '_').title() + "__generated_j2cl_test_input"
  java_code = _TEMPLATE % (PACKAGE_NAME, name, java_package, test_class, java_class)

  native.genrule(
      name = java_class,
      outs = [java_class + ".java"],
      cmd="echo \"%s\" > $@" % java_code,
      tags=["manual", "notap"],
      testonly = 1,
  )

  return java_class + ".java"

def j2cl_generate_jsunit_suite(name, test_class, deps, tags = []):
  """Macro for cross compiling a JUnit Suite to JavaScript testSuite"""

  test_input = _generate_test_input(name, test_class)
  # This target triggers our Java annotation processor and generates the
  # plumbing between a jsunit_test and transpiled JUnit tests
  # It's outputs are:
  #  - test_summary.json that lists all jsunit test suites
  #  - A .testsuite file for every JUnit test containing a jsunit tests suite
  #  - A Java class (JsType) that contains the bridge between the jsunit test
  #        suite and the JUnit test. This is being used from the .testsuite
  #        file.
  #
  # We separated this out into a separate target so we can have dependencies
  # here that users might not have in their tests (e.g. jsinterop annotations).
  # We need the extra dep here on user provided dependencies since our generated
  # code refers to user written code (test cases).
  # Note that test suites are generated with .testsuite extension to avoid
  # j2cl_library automically including them as source.
  j2cl_library(
      name = name + "_lib",
      srcs = [test_input],
      deps = deps +  [
          "//internal_do_not_use:internal_junit_annotations",
          "//third_party/java/gwt:gwt-jsinterop-annotations-j2cl",
      ],
      _js_deps = ["//javascript/closure/testing:testcase"],
      plugins = ["//internal_do_not_use:junit_processor_plugin"],
      testonly = 1,
      tags = tags,
      generate_build_test=False,
  )

  # The Java annotation processor on the above target generates jsunit suites
  # (.testsuite files), but the same jar file also contains unrelated things
  # (e.g. class files), this genrules takes the jar file as input and creates
  # a new zip file that only contains the generated javascript (.testsuite
  # renamed to .js) and the test_summary.json file.
  # This is the format that jsunit_test will later expect.
  # TODO(goktug): use j2cl_library directly from jsunit_test instead of
  # extracting files from jar (output js zip can include all the required
  # files.)
  out_jar = ":lib" + name + "_lib_java_library.jar"
  native.genrule(
      name=name + "_transpile_gen",
      outs=[name + ".js.zip"],
      cmd="\n".join([
          "unzip -q $(location %s) *.testsuite *.json -d zip_out/" % out_jar,
          "cd zip_out/",
          "for f in $$(find -name *.testsuite); do mv $$f $${f/.testsuite/.js}; done",
          "zip -q -r ../$@ .",
      ]),
      testonly=1,
      tags=["manual", "notap"],
      tools=[out_jar],
  )
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Datastore manages file retrieval for PUDL datasets."""

import argparse
import hashlib
import json
import logging
import re
import sys
from pathlib import Path

import requests
import yaml

# The Zenodo tokens recorded here should have read-only access to our archives.
# Including them here is correct in order to allow public use of this tool, so
# long as we stick to read-only keys.

TOKEN = {
    "sandbox": "WX1FW2JOrkcCvoGutoZaGS69W5P9A73wybsToLeo4nt6NP3moeIg7sPBw8nI",
    "production": "N/A"
}

DOI = {
    "sandbox": {
        "eia860": "10.5072/zenodo.504556",
        "eia861": "10.5072/zenodo.504558",
        "eia923": "10.5072/zenodo.504560",
        "epaipm": "10.5072/zenodo.602953",
        "epacems": "10.5072/zenodo.638878",
        "ferc1": "10.5072/zenodo.504562"
    },
    "production": {}
}

PUDL_YML = Path.home() / ".pudl.yml"

class Datastore:
    """Handle connections and downloading of Zenodo Source archives."""

    def __init__(self, pudl_in, loglevel="WARNING", verbose=False, sandbox=False):
        """
        Datastore manages file retrieval for PUDL datasets.

        It pulls files from Zenodo archives as needed, and caches them
        locally. If the environment variable PUDL_IN is set, the datastore
        is assumed to be in PUDL_IN / data. Otherwise PUDL_IN is read
        from the user's ~/.pudl.yml settings file.

        Args:
            pudl_in (Path): path to the root pudl data directory
            loglevel: str, logging level
            verbose: boolean. If true, logs printed to stdout
            sandbox: boolean. If true, use the sandbox server instead of
                     production
        """
        self.pudl_in = pudl_in
        logger = logging.Logger(__name__)
        logger.setLevel(loglevel)

        if verbose:
            logger.addHandler(logging.StreamHandler())

        self.logger = logger
        logger.info("Logging at loglevel %s" % loglevel)

        if sandbox:
            self._dois = DOI["sandbox"]
            self.token = TOKEN["sandbox"]
            self.api_root = "https://sandbox.zenodo.org/api"

        else:
            self._dois = DOI["production"]
            self.token = TOKEN["production"]
            self.api_root = "https://zenodo.org/api"

    # Location conversion & helpers

    def doi(self, dataset):
        """
        Produce the DOI for a given dataset, or log & raise error.

        Args:
            datapackage: the name of the datapackage

        Returns:
            str, doi for the datapackage

        """
        try:
            return self._dois[dataset]
        except KeyError:
            msg = "No DOI available for %s" % dataset
            self.logger.error(msg)
            raise ValueError(msg)

    def doi_to_url(self, doi):
        """
        Given a DOI, produce the API url to retrieve it.

        Args:
            doi: str, the doi (concept doi) to retrieve, per
                 https://help.zenodo.org/

        Returns:
            url to get the deposition from the api
        """
        match = re.search(r"zenodo.([\d]+)", doi)

        if match is None:
            msg = "Invalid doi %s" % doi
            self.logger.error(msg)
            raise ValueError(msg)

        zen_id = int(match.groups()[0])
        return "%s/deposit/depositions/%d" % (self.api_root, zen_id)

    def local_path(self, dataset, filename=None):
        """
        Produce the local path for a given dataset.

        Args:
            dataset: the name of the dataset
            filename: optional filename as it would appear, if included in the
                      path

        Return:
            str: a path
        """
        doi_dirname = re.sub("/", "-", self.doi(dataset))
        directory = self.pudl_in / "data" / dataset / doi_dirname

        if filename is None:
            return directory

        return directory / filename

    def save_datapackage_json(self, dataset, dpkg):
        """
        Save a datapackage.json file.  Overwrite any previous version.

        Args:
            dataset (str): name of the dataset, as available in DOI
            dpkg (dict): dict matching frictionless datapackage spec

        Returns:
            Path of the saved datapackage
        """
        path = self.local_path(dataset, filename="datapackage.json")
        path.parent.mkdir(parents=True, exist_ok=True)
        text = json.dumps(dpkg, sort_keys=True)

        with path.open("w") as f:
            f.write(text)
            self.logger.debug("%s saved.", path)

        return path

    # Datapackage metadata

    def remote_datapackage_json(self, doi):
        """
        Produce the contents of a remote datapackage.json.

        Args:
            doi: the DOI

        Returns:
            dict representation of the datapackage.json file as available on
            Zenodo, or raises an error
        """
        dpkg_url = self.doi_to_url(doi)
        response = requests.get(
            dpkg_url, params={"access_token": self.token})

        if response.status_code > 299:
            msg = "Failed to retrieve %s" % dpkg_url
            self.logger.error(msg)
            raise ValueError(msg)

        jsr = response.json()
        files = {x["filename"]: x for x in jsr["files"]}

        response = requests.get(
            files["datapackage.json"]["links"]["download"],
            params={"access_token": self.token})

        if response.status_code > 299:
            msg = "Failed to retrieve datapackage for %s: %s" % (
                doi, response.text)
            self.logger.error(msg)
            raise ValueError(msg)

        return json.loads(response.text)

    def datapackage_json(self, dataset, force_download=False):
        """
        Produce the contents of datapackage.json, cache as necessary.

        Args:
            dataset (str): name of the dataset, must be available in the DOIS

        Return:
            dict representation of the datapackage.json
        """
        path = self.local_path(dataset, filename="datapackage.json")

        if force_download or not path.exists():
            doi = self.doi(dataset)
            dpkg = self.remote_datapackage_json(doi)
            self.save_datapackage_json(dataset, dpkg)

        with path.open("r") as f:
            return yaml.safe_load(f)

    # Remote resource retrieval

    def is_remote(self, resource):
        """
        Determine whether a described resource is located on a remote server.

        Args:
            resource: dict, a resource descriptor from a frictionless data
                      package
        Returns:
            boolean
        """
        return resource["path"][:8] == "https://" or \
            resource["path"][:7] == "http://"

    def passes_filters(self, resource, filters):
        """
        Test whether file metadata passes given filters.

        Args:
            resource: dict, a "resource" descriptior from a frictionless
                      datapackage
            filters: dict, pairs that must match in order for a
                     resource to pass
        Returns:
            boolean, True if the resource parts pass the filters
        """
        for key, _ in filters.items():

            part_val = resource["parts"].get(key, None)

            if part_val != filters[key]:
                self.logger.debug(
                    "Filtered %s on %s: %s != %s", resource["name"], key,
                    part_val, filters[key])
                return False

        return True

    def download_resource(self, resource, directory):
        """
        Download a frictionless datapackage resource.

        Args:
            resource: dict, a remotely located resource descriptior from a
                      frictionless datapackage
            directory: the directory where the resource should be saved
        Returns:
            Path of the saved resource, or none on failure
        """
        response = requests.get(
            resource["path"], params={"access_token": self.token})

        if response.status_code >= 299:
            msg = "Failed to download %s, %s" % (
                resource["path"], response.text)
            self.logger.error(msg)
            raise RuntimeError(msg)

        local_path = directory / resource["name"]

        with local_path.open("wb") as f:
            f.write(response.content)
            self.logger.debug("Cached %s" % local_path)

        return local_path

    def _validate_dataset(self, dataset):
        """
        Make sure each resource on disc has the right md5sum.

        Args:
            dataset: name of a dataset
        Returns:
            Bool, True if local resources appear to be correct.
        """
        dpkg = self.datapackage_json(dataset)
        ok = True

        for r in dpkg["resources"]:

            if self.is_remote(r):
                self.logger.debug("%s not cached, skipping validation"
                                  % r["path"])
                continue

            path = Path(r["path"])
            with path.open("rb") as f:
                m = hashlib.md5()  # nosec
                m.update(f.read())

            if m.hexdigest() == r["hash"]:
                self.logger.debug("%s appears valid" % path)
            else:
                self.logger.warning("%s md5 mismatch" % path)
                ok = False

        return ok

    def validate(self, dataset=None):
        """
        Validate all datasets, or just the specified one.

        Args:
            dataset: name of a dataset.
        Returns:
            Bool, True if local resources appear to be correct.
        """
        if dataset is not None:
            return self._validate_dataset(dataset)

        valid = True

        for dataset in self._dois.keys():
            valid = valid and self._validate_dataset(dataset)

        return valid

    def get_resources(self, dataset, **kwargs):
        """
        Produce resource descriptors as requested.

        Any resource listed as remote will be downloaded and the
        datapackage.json will be updated.

        Args:
            dataset (str): name of the dataset, must be available in the DOIS
            **kwargs: limit retrieved files to those where the
                datapackage.json["parts"] key & val pairs match provided
                keywords. Eg. year=2011 or state="md"
        Returns:
            list of dicts, each representing a resource per the frictionless
            datapackage spec. For a given r, Path(r["path"]) should open the
            local file, r["parts"] should provide metadata identifiers.
        """
        filters = dict(**kwargs)

        dpkg = self.datapackage_json(dataset)
        self.logger.debug("Datapackage lists %d resources" %
                          len(dpkg["resources"]))

        for r in dpkg["resources"]:

            if self.passes_filters(r, filters):

                if self.is_remote(r):
                    local = self.download_resource(r, self.local_path(dataset))
                    r["path"] = str(local)
                    self.save_datapackage_json(dataset, dpkg)

                yield r

def main_arguments():
    """Collect the command line arguments."""
    prod_dois = "\n".join(["    - %s" % x for x in DOI["production"].keys()])
    sand_dois = "\n".join(["    - %s" % x for x in DOI["sandbox"].keys()])

    dataset_msg = "--Available datasets--\n \nProduction:\n%s\n \nSandbox:\n%s" \
        % (prod_dois, sand_dois)

    parser = argparse.ArgumentParser(
        description="Download and cache ETL source data from Zenodo.",
        epilog=dataset_msg,
        formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument(
        "--dataset", help="Get only a specified dataset. Default gets all.")
    parser.add_argument(
        "--pudl_in",
        help="Override pudl_in directory, defaults to option .pudl.yml")
    parser.add_argument(
        "--validate", help="Validate existing cache.", const=True,
        default=False, action="store_const")
    parser.add_argument(
        "--sandbox", help="Use sandbox server instead of production.",
        action="store_const", const=True, default=False)
    parser.add_argument(
        "--loglevel", help="Set logging level", default="WARNING")
    parser.add_argument(
        "--verbose", help="Display logging messages", default=False,
        action="store_const", const=True)

    return parser.parse_args()

def main():
    """Cache datasets."""
    args = main_arguments()
    dataset = getattr(args, "dataset", None)
    pudl_in = getattr(args, "pudl_in", None)

    if pudl_in is None:
        with PUDL_YML.open() as f:
            cfg = yaml.safe_load(f)
            pudl_in = Path(cfg["pudl_in"])
    else:
        pudl_in = Path(pudl_in)

    ds = Datastore(pudl_in, loglevel=args.loglevel, verbose=args.verbose,
                   sandbox=args.sandbox)

    if args.validate:
        ds.validate(dataset)
        return

    list(ds.get_resources(dataset))

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Datastore bugfix<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Datastore manages file retrieval for PUDL datasets."""

import argparse
import hashlib
import json
import logging
import re
import sys
from pathlib import Path

import requests
import yaml

# The Zenodo tokens recorded here should have read-only access to our archives.
# Including them here is correct in order to allow public use of this tool, so
# long as we stick to read-only keys.

TOKEN = {
    "sandbox": "WX1FW2JOrkcCvoGutoZaGS69W5P9A73wybsToLeo4nt6NP3moeIg7sPBw8nI",
    "production": "N/A"
}

DOI = {
    "sandbox": {
        "eia860": "10.5072/zenodo.504556",
        "eia861": "10.5072/zenodo.504558",
        "eia923": "10.5072/zenodo.504560",
        "epaipm": "10.5072/zenodo.602953",
        "epacems": "10.5072/zenodo.638878",
        "ferc1": "10.5072/zenodo.504562"
    },
    "production": {}
}

PUDL_YML = Path.home() / ".pudl.yml"

class Datastore:
    """Handle connections and downloading of Zenodo Source archives."""

    def __init__(self, pudl_in, loglevel="WARNING", verbose=False, sandbox=False):
        """
        Datastore manages file retrieval for PUDL datasets.

        It pulls files from Zenodo archives as needed, and caches them
        locally. If the environment variable PUDL_IN is set, the datastore
        is assumed to be in PUDL_IN / data. Otherwise PUDL_IN is read
        from the user's ~/.pudl.yml settings file.

        Args:
            pudl_in (Path): path to the root pudl data directory
            loglevel: str, logging level
            verbose: boolean. If true, logs printed to stdout
            sandbox: boolean. If true, use the sandbox server instead of
                     production
        """
        self.pudl_in = pudl_in
        logger = logging.Logger(__name__)
        logger.setLevel(loglevel)

        if verbose:
            logger.addHandler(logging.StreamHandler())

        self.logger = logger
        logger.info("Logging at loglevel %s" % loglevel)

        if sandbox:
            self._dois = DOI["sandbox"]
            self.token = TOKEN["sandbox"]
            self.api_root = "https://sandbox.zenodo.org/api"

        else:
            raise NotImplementedError(
                "Production archive not ready. Use --sandbox.")
            self._dois = DOI["production"]
            self.token = TOKEN["production"]
            self.api_root = "https://zenodo.org/api"

    # Location conversion & helpers

    def doi(self, dataset):
        """
        Produce the DOI for a given dataset, or log & raise error.

        Args:
            datapackage: the name of the datapackage

        Returns:
            str, doi for the datapackage

        """
        try:
            return self._dois[dataset]
        except KeyError:
            msg = "No DOI available for %s" % dataset
            self.logger.error(msg)
            raise ValueError(msg)

    def doi_to_url(self, doi):
        """
        Given a DOI, produce the API url to retrieve it.

        Args:
            doi: str, the doi (concept doi) to retrieve, per
                 https://help.zenodo.org/

        Returns:
            url to get the deposition from the api
        """
        match = re.search(r"zenodo.([\d]+)", doi)

        if match is None:
            msg = "Invalid doi %s" % doi
            self.logger.error(msg)
            raise ValueError(msg)

        zen_id = int(match.groups()[0])
        return "%s/deposit/depositions/%d" % (self.api_root, zen_id)

    def local_path(self, dataset, filename=None):
        """
        Produce the local path for a given dataset.

        Args:
            dataset: the name of the dataset
            filename: optional filename as it would appear, if included in the
                      path

        Return:
            str: a path
        """
        doi_dirname = re.sub("/", "-", self.doi(dataset))
        directory = self.pudl_in / "data" / dataset / doi_dirname

        if filename is None:
            return directory

        return directory / filename

    def save_datapackage_json(self, dataset, dpkg):
        """
        Save a datapackage.json file.  Overwrite any previous version.

        Args:
            dataset (str): name of the dataset, as available in DOI
            dpkg (dict): dict matching frictionless datapackage spec

        Returns:
            Path of the saved datapackage
        """
        path = self.local_path(dataset, filename="datapackage.json")
        path.parent.mkdir(parents=True, exist_ok=True)
        text = json.dumps(dpkg, sort_keys=True)

        with path.open("w") as f:
            f.write(text)
            self.logger.debug("%s saved.", path)

        return path

    # Datapackage metadata

    def remote_datapackage_json(self, doi):
        """
        Produce the contents of a remote datapackage.json.

        Args:
            doi: the DOI

        Returns:
            dict representation of the datapackage.json file as available on
            Zenodo, or raises an error
        """
        dpkg_url = self.doi_to_url(doi)
        response = requests.get(
            dpkg_url, params={"access_token": self.token})

        if response.status_code > 299:
            msg = "Failed to retrieve %s" % dpkg_url
            self.logger.error(msg)
            raise ValueError(msg)

        jsr = response.json()
        files = {x["filename"]: x for x in jsr["files"]}

        response = requests.get(
            files["datapackage.json"]["links"]["download"],
            params={"access_token": self.token})

        if response.status_code > 299:
            msg = "Failed to retrieve datapackage for %s: %s" % (
                doi, response.text)
            self.logger.error(msg)
            raise ValueError(msg)

        return json.loads(response.text)

    def datapackage_json(self, dataset, force_download=False):
        """
        Produce the contents of datapackage.json, cache as necessary.

        Args:
            dataset (str): name of the dataset, must be available in the DOIS

        Return:
            dict representation of the datapackage.json
        """
        path = self.local_path(dataset, filename="datapackage.json")

        if force_download or not path.exists():
            doi = self.doi(dataset)
            dpkg = self.remote_datapackage_json(doi)
            self.save_datapackage_json(dataset, dpkg)

        with path.open("r") as f:
            return yaml.safe_load(f)

    # Remote resource retrieval

    def is_remote(self, resource):
        """
        Determine whether a described resource is located on a remote server.

        Args:
            resource: dict, a resource descriptor from a frictionless data
                      package
        Returns:
            boolean
        """
        return resource["path"][:8] == "https://" or \
            resource["path"][:7] == "http://"

    def passes_filters(self, resource, filters):
        """
        Test whether file metadata passes given filters.

        Args:
            resource: dict, a "resource" descriptior from a frictionless
                      datapackage
            filters: dict, pairs that must match in order for a
                     resource to pass
        Returns:
            boolean, True if the resource parts pass the filters
        """
        for key, _ in filters.items():

            part_val = resource["parts"].get(key, None)

            if part_val != filters[key]:
                self.logger.debug(
                    "Filtered %s on %s: %s != %s", resource["name"], key,
                    part_val, filters[key])
                return False

        return True

    def download_resource(self, resource, directory):
        """
        Download a frictionless datapackage resource.

        Args:
            resource: dict, a remotely located resource descriptior from a
                      frictionless datapackage
            directory: the directory where the resource should be saved
        Returns:
            Path of the saved resource, or none on failure
        """
        response = requests.get(
            resource["path"], params={"access_token": self.token})

        if response.status_code >= 299:
            msg = "Failed to download %s, %s" % (
                resource["path"], response.text)
            self.logger.error(msg)
            raise RuntimeError(msg)

        local_path = directory / resource["name"]

        with local_path.open("wb") as f:
            f.write(response.content)
            self.logger.debug("Cached %s" % local_path)

        return local_path

    def _validate_dataset(self, dataset):
        """
        Make sure each resource on disc has the right md5sum.

        Args:
            dataset: name of a dataset
        Returns:
            Bool, True if local resources appear to be correct.
        """
        dpkg = self.datapackage_json(dataset)
        ok = True

        for r in dpkg["resources"]:

            if self.is_remote(r):
                self.logger.debug("%s not cached, skipping validation"
                                  % r["path"])
                continue

            path = Path(r["path"])
            with path.open("rb") as f:
                m = hashlib.md5()  # nosec
                m.update(f.read())

            if m.hexdigest() == r["hash"]:
                self.logger.debug("%s appears valid" % path)
            else:
                self.logger.warning("%s md5 mismatch" % path)
                ok = False

        return ok

    def validate(self, dataset=None):
        """
        Validate all datasets, or just the specified one.

        Args:
            dataset: name of a dataset.
        Returns:
            Bool, True if local resources appear to be correct.
        """
        if dataset is not None:
            return self._validate_dataset(dataset)

        valid = True

        for dataset in self._dois.keys():
            valid = valid and self._validate_dataset(dataset)

        return valid

    def get_resources(self, dataset, **kwargs):
        """
        Produce resource descriptors as requested.

        Any resource listed as remote will be downloaded and the
        datapackage.json will be updated.

        Args:
            dataset (str): name of the dataset, must be available in the DOIS
            **kwargs: limit retrieved files to those where the
                datapackage.json["parts"] key & val pairs match provided
                keywords. Eg. year=2011 or state="md"
        Returns:
            list of dicts, each representing a resource per the frictionless
            datapackage spec. For a given r, Path(r["path"]) should open the
            local file, r["parts"] should provide metadata identifiers.
        """
        filters = dict(**kwargs)

        dpkg = self.datapackage_json(dataset)
        self.logger.debug("Datapackage lists %d resources" %
                          len(dpkg["resources"]))

        for r in dpkg["resources"]:

            if self.passes_filters(r, filters):

                if self.is_remote(r):
                    local = self.download_resource(r, self.local_path(dataset))
                    r["path"] = str(local)
                    self.save_datapackage_json(dataset, dpkg)

                yield r

def main_arguments():
    """Collect the command line arguments."""
    prod_dois = "\n".join(["    - %s" % x for x in DOI["production"].keys()])
    sand_dois = "\n".join(["    - %s" % x for x in DOI["sandbox"].keys()])

    dataset_msg = "--Available datasets--\n \nProduction:\n%s\n \nSandbox:\n%s" \
        % (prod_dois, sand_dois)

    parser = argparse.ArgumentParser(
        description="Download and cache ETL source data from Zenodo.",
        epilog=dataset_msg,
        formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument(
        "--dataset", help="Get only a specified dataset. Default gets all.")
    parser.add_argument(
        "--pudl_in",
        help="Override pudl_in directory, defaults to option .pudl.yml")
    parser.add_argument(
        "--validate", help="Validate existing cache.", const=True,
        default=False, action="store_const")
    parser.add_argument(
        "--sandbox", help="Use sandbox server instead of production.",
        action="store_const", const=True, default=False)
    parser.add_argument(
        "--loglevel", help="Set logging level", default="WARNING")
    parser.add_argument(
        "--verbose", help="Display logging messages", default=False,
        action="store_const", const=True)

    return parser.parse_args()

def main():
    """Cache datasets."""
    args = main_arguments()
    dataset = getattr(args, "dataset", None)
    pudl_in = getattr(args, "pudl_in", None)

    if pudl_in is None:
        with PUDL_YML.open() as f:
            cfg = yaml.safe_load(f)
            pudl_in = Path(cfg["pudl_in"])
    else:
        pudl_in = Path(pudl_in)

    ds = Datastore(pudl_in, loglevel=args.loglevel, verbose=args.verbose,
                   sandbox=args.sandbox)

    if dataset is None:
        if args.sandbox:
            datasets = DOI["sandbox"].keys()
        else:
            datasets = DOI["production"].keys()
    else:
        datasets = [dataset]

    for selection in datasets:

        if args.validate:
            ds.validate(selection)
            continue

        list(ds.get_resources(selection))

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


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

from aiohttp.hdrs import (
    METH_GET,
    METH_POST,
    METH_PUT,
    METH_DELETE
)
from aiohttp import web, MultiDictProxy
from djaio.core.exceptions import BadRequestException
from djaio.core.models import NullInput, NullOutput

from djaio.core.utils import get_int_or_none
from schematics.exceptions import ModelConversionError, ConversionError, DataError, ValidationError

class BaseMethod(object):
    def __init__(self, input_model=NullInput, output_model=NullOutput, description=None):
        self.result = None
        self.input_model = input_model
        self.output_model = output_model
        self.total = None
        self.success = None
        self.errors = None
        self.params = None
        self.output = None
        self.pagination = None
        self.limit = None
        self.offset = None
        self.app = None
        self.settings = None
        self.description = description

    def process_request(self, multi):
        #Override it for your purposes
        params = {}
        # Here we convert a MultiDict to simple python dict.py
        if isinstance(multi, MultiDictProxy):
            for k in set(multi.keys()):
                v = multi.getall(k)
                params[k] = v if len(v) > 1 else v[0]
        elif isinstance(multi, dict):
            params = multi
        else:
            raise BadRequestException(message='request params must be a dict-like object')
        return params

    async def from_http(self, request):
        self.total = None
        self.success = None
        self.errors = []
        self.params = None
        self.output = None
        self.pagination = None
        self.limit = None
        self.offset = None
        if not isinstance(request, web.Request):
            raise web.HTTPBadRequest()
        try:
            req_params = {}
            # if GET or DELETE we read a query params
            if request.method in (METH_GET, METH_DELETE):
                req_params = self.process_request(request.GET)
            # else we read a POST-data
            elif request.method in (METH_PUT, METH_POST):
                try:
                    req_params = self.process_request(await request.json())
                except (ValueError, TypeError):
                    req_params = self.process_request(await request.post())

            # Here we add or owerride params by PATH-params.
            # If it exist
            if request.match_info:
                req_params.update(request.match_info.copy())

            self.limit = get_int_or_none(request.headers.get('X-Limit')) or \
                         get_int_or_none(req_params.pop('limit', None)) or \
                         get_int_or_none(request.app.settings.LIMIT)
            self.offset = get_int_or_none(request.headers.get('X-Offset')) or \
                          get_int_or_none(req_params.pop('offset', None)) or \
                          get_int_or_none(request.app.settings.OFFSET)

            params = self.input_model(req_params, strict=False)
            params.validate()
            self.params = params.to_primitive()
        except (ModelConversionError, ConversionError, DataError) as exc:
            errors = []
            if isinstance(exc.messages, list):
                errors = [x.summary for x in exc.messages]
            else:
                for k, v in exc.messages.items():
                    sub_errors = {}
                    if isinstance(v, dict):
                        for sub_k, error in v.items():
                            if isinstance(error, ConversionError) or isinstance(error, ValidationError):
                                sub_errors[sub_k] = [x.summary for x in error.messages]
                        errors.append({k: sub_errors})

                    elif isinstance(v, ConversionError) or isinstance(v, ValidationError):
                        errors.append({k: [x.summary for x in v.messages]})

                    elif isinstance(v, list):
                        errors.append({k: [x.summary for x in v]})
                    elif isinstance(v, str):
                        errors.append({k: v})
            raise BadRequestException(message=errors)

        self.result = []
        self.app = request.app
        self.settings = request.app.settings

    async def execute(self):
        raise NotImplementedError('Please override `execute()` method.')

    def get_pagination(self):
        _pagination_object = {
            'total': self.total or len(self.result),
            'limit': self.limit,
            'offset': self.offset
        }
        self.pagination = _pagination_object
        return self.pagination

    async def get_output(self):
        self.result = await self.execute()
        self.output = {
            'success': not self.errors
        }
        if not self.errors:
            if type(self.result) in (list, tuple) or isinstance(self.result, map):
                self.output['result'] = [self.output_model(x, strict=False).to_primitive() for x in self.result]
            elif type(self.result) == dict and self.result:
                self.output['result'] = self.output_model(self.result, strict=False).to_primitive()
            else:
                self.output['result'] = self.result
        else:
            self.output.update({'errors': self.errors})
        pagination = self.get_pagination()
        if pagination:
            self.output.update({'pagination': pagination})
        return self.output
<commit_msg>set total to 0 if nothing found<commit_after># -*- coding: utf-8 -*-
import json

from aiohttp.hdrs import (
    METH_GET,
    METH_POST,
    METH_PUT,
    METH_DELETE
)
from aiohttp import web, MultiDictProxy
from djaio.core.exceptions import BadRequestException
from djaio.core.models import NullInput, NullOutput

from djaio.core.utils import get_int_or_none
from schematics.exceptions import ModelConversionError, ConversionError, DataError, ValidationError

class BaseMethod(object):
    def __init__(self, input_model=NullInput, output_model=NullOutput, description=None):
        self.result = None
        self.input_model = input_model
        self.output_model = output_model
        self.total = None
        self.success = None
        self.errors = None
        self.params = None
        self.output = None
        self.pagination = None
        self.limit = None
        self.offset = None
        self.app = None
        self.settings = None
        self.description = description

    def process_request(self, multi):
        #Override it for your purposes
        params = {}
        # Here we convert a MultiDict to simple python dict.py
        if isinstance(multi, MultiDictProxy):
            for k in set(multi.keys()):
                v = multi.getall(k)
                params[k] = v if len(v) > 1 else v[0]
        elif isinstance(multi, dict):
            params = multi
        else:
            raise BadRequestException(message='request params must be a dict-like object')
        return params

    async def from_http(self, request):
        self.total = None
        self.success = None
        self.errors = []
        self.params = None
        self.output = None
        self.pagination = None
        self.limit = None
        self.offset = None
        if not isinstance(request, web.Request):
            raise web.HTTPBadRequest()
        try:
            req_params = {}
            # if GET or DELETE we read a query params
            if request.method in (METH_GET, METH_DELETE):
                req_params = self.process_request(request.GET)
            # else we read a POST-data
            elif request.method in (METH_PUT, METH_POST):
                try:
                    req_params = self.process_request(await request.json())
                except (ValueError, TypeError):
                    req_params = self.process_request(await request.post())

            # Here we add or owerride params by PATH-params.
            # If it exist
            if request.match_info:
                req_params.update(request.match_info.copy())

            self.limit = get_int_or_none(request.headers.get('X-Limit')) or \
                         get_int_or_none(req_params.pop('limit', None)) or \
                         get_int_or_none(request.app.settings.LIMIT)
            self.offset = get_int_or_none(request.headers.get('X-Offset')) or \
                          get_int_or_none(req_params.pop('offset', None)) or \
                          get_int_or_none(request.app.settings.OFFSET)

            params = self.input_model(req_params, strict=False)
            params.validate()
            self.params = params.to_primitive()
        except (ModelConversionError, ConversionError, DataError) as exc:
            errors = []
            if isinstance(exc.messages, list):
                errors = [x.summary for x in exc.messages]
            else:
                for k, v in exc.messages.items():
                    sub_errors = {}
                    if isinstance(v, dict):
                        for sub_k, error in v.items():
                            if isinstance(error, ConversionError) or isinstance(error, ValidationError):
                                sub_errors[sub_k] = [x.summary for x in error.messages]
                        errors.append({k: sub_errors})

                    elif isinstance(v, ConversionError) or isinstance(v, ValidationError):
                        errors.append({k: [x.summary for x in v.messages]})

                    elif isinstance(v, list):
                        errors.append({k: [x.summary for x in v]})
                    elif isinstance(v, str):
                        errors.append({k: v})
            raise BadRequestException(message=errors)

        self.result = []
        self.app = request.app
        self.settings = request.app.settings

    async def execute(self):
        raise NotImplementedError('Please override `execute()` method.')

    def get_pagination(self):
        total = None
        if self.result is None:
            total =0
        elif isinstance(self.result, list):
            total = len(self.result)
        _pagination_object = {
            'total': self.total or total,
            'limit': self.limit,
            'offset': self.offset
        }
        self.pagination = _pagination_object
        return self.pagination

    async def get_output(self):
        self.result = await self.execute()
        self.output = {
            'success': not self.errors
        }
        if not self.errors:
            if type(self.result) in (list, tuple) or isinstance(self.result, map):
                self.output['result'] = [self.output_model(x, strict=False).to_primitive() for x in self.result]
            elif type(self.result) == dict and self.result:
                self.output['result'] = self.output_model(self.result, strict=False).to_primitive()
            else:
                self.output['result'] = self.result
        else:
            self.output.update({'errors': self.errors})
        pagination = self.get_pagination()
        if pagination:
            self.output.update({'pagination': pagination})
        return self.output
<|endoftext|>


<commit_before>__all__ = ("VirtualHeap",
           "SizedVirtualHeap")

import os
import sys
import subprocess
import random
import string
import tempfile

from six.moves import range

from pyoram.util._virtual_heap_helper import lib as _clib
from pyoram.util.misc import log2floor

numerals = ''.join([c for c in string.printable \
                  if ((c not in string.whitespace) and \
                      (c != '+') and (c != '-') and \
                      (c != '"') and (c != "'") and \
                      (c != '\\') and (c != '/'))])
numeral_index = dict((c,i) for i,c in enumerate(numerals))

# The maximum heap base for which base k labels
# can be produced.
max_k_labeled = len(numerals)

def base10_integer_to_basek_string(k, x):
    """Convert an integer into a base k string."""
    if not (2 <= k <= max_k_labeled):
        raise ValueError("k must be in range [2, %d]: %s"
                         % (max_k_labeled, k))
    return ((x == 0) and numerals[0]) or \
        (base10_integer_to_basek_string(k, x // k).\
         lstrip(numerals[0]) + numerals[x % k])

def basek_string_to_base10_integer(k, x):
    """Convert a base k string into an integer."""
    assert 1 < k <= max_k_labeled
    return sum(numeral_index[c]*(k**i)
               for i, c in enumerate(reversed(x)))

# _clib defines a faster version of this function
def calculate_bucket_level(k, b):
    """
    Calculate the level in which a 0-based bucket
    lives inside of a k-ary heap.
    """
    assert k >= 2
    if k == 2:
        return log2floor(b+1)
    v = (k - 1) * (b + 1) + 1
    h = 0
    while k**(h+1) < v:
        h += 1
    return h

# _clib defines a faster version of this function
def calculate_last_common_level(k, b1, b2):
    """
    Calculate the highest level after which the
    paths from the root to these buckets diverge.
    """
    l1 = calculate_bucket_level(k, b1)
    l2 = calculate_bucket_level(k, b2)
    while l1 > l2:
        b1 = (b1-1)//k
        l1 -= 1
    while l2 > l1:
        b2 = (b2-1)//k
        l2 -= 1
    while b1 != b2:
        b1 = (b1-1)//k
        b2 = (b2-1)//k
        l1 -= 1
    return l1

def calculate_necessary_heap_height(k, n):
    """
    Calculate the necessary k-ary heap height
    to store n buckets.
    """
    assert n >= 1
    return calculate_bucket_level(k, n-1)

def calculate_bucket_count_in_heap_with_height(k, h):
    """
    Calculate the number of buckets in a
    k-ary heap of height h.
    """
    assert h >= 0
    return ((k**(h+1)) - 1) // (k - 1)

def calculate_bucket_count_in_heap_at_level(k, l):
    """
    Calculate the number of buckets in a
    k-ary heap at level l.
    """
    assert l >= 0
    return k**l

def calculate_leaf_bucket_count_in_heap_with_height(k, h):
    """
    Calculate the number of buckets in the
    leaf-level of a k-ary heap of height h.
    """
    return calculate_bucket_count_in_heap_at_level(k, h)

def create_node_type(k):

    class VirtualHeapNode(object):
        __slots__ = ("bucket", "level")
        def __init__(self, bucket):
            assert bucket >= 0
            self.bucket = bucket
            self.level = _clib.calculate_bucket_level(self.k, self.bucket)

        def __hash__(self):
            return self.bucket.__hash__()
        def __lt__(self, other):
            return self.bucket < other
        def __le__(self, other):
            return self.bucket <= other
        def __eq__(self, other):
            return self.bucket == other
        def __ne__(self, other):
            return self.bucket != other
        def __gt__(self, other):
            return self.bucket > other
        def __ge__(self, other):
            return self.bucket >= other
        def last_common_level(self, n):
            return _clib.calculate_last_common_level(self.k,
                                                     self.bucket,
                                                     n.bucket)
        def child_node(self, c):
            assert type(c) is int
            assert 0 <= c < self.k
            return VirtualHeapNode(self.k * self.bucket + 1 + c)
        def parent_node(self):
            if self.bucket != 0:
                return VirtualHeapNode((self.bucket - 1)//self.k)
            return None
        def ancestor_node_at_level(self, level):
            if level > self.level:
                return None
            current = self
            while current.level != level:
                current = current.parent_node()
            return current
        def bucket_path_to_root(self):
            bucket = self.bucket
            yield bucket
            while bucket != 0:
                bucket = (bucket - 1)//self.k
                yield bucket
        def bucket_path_from_root(self):
            return list(reversed(list(self.bucket_path_to_root())))

        #
        # Expensive Functions
        #
        def __repr__(self):
            try:
                label = self.label()
            except ValueError:
                # presumably, k is too large
                label = "<unknown>"
            return ("VirtualHeapNode(k=%s, bucket=%s, level=%s, label=%r)"
                    % (self.k, self.bucket, self.level, label))
        def __str__(self):
            """Returns a tuple (<level>, <bucket offset within level>)."""
            if self.bucket != 0:
                return ("(%s, %s)"
                        % (self.level,
                           self.bucket -
                           calculate_bucket_count_in_heap_with_height(self.k,
                                                                self.level-1)))
            assert self.level == 0
            return "(0, 0)"

        def label(self):
            assert 0 <= self.bucket
            if self.level == 0:
                return ''
            b_offset = self.bucket - \
                       calculate_bucket_count_in_heap_with_height(self.k,
                                                            self.level-1)
            basek = base10_integer_to_basek_string(self.k, b_offset)
            return basek.zfill(self.level)

        def is_node_on_path(self, n):
            if n.level <= self.level:
                n_label = n.label()
                if n_label == "":
                    return True
                return self.label().startswith(n_label)
            return False

    VirtualHeapNode.k = k

    return VirtualHeapNode

class VirtualHeap(object):

    clib = _clib
    random = random.SystemRandom()

    def __init__(self, k, blocks_per_bucket=1):
        assert 1 < k
        assert blocks_per_bucket >= 1
        self._k = k
        self._blocks_per_bucket = blocks_per_bucket
        self.Node = create_node_type(k)

    @property
    def k(self):
        return self._k

    def node_label_to_bucket(self, label):
        if len(label) > 0:
            return \
                (calculate_bucket_count_in_heap_with_height(self.k,
                                                      len(label)-1) +
                 basek_string_to_base10_integer(self.k, label))
        return 0

    #
    # Buckets (0-based integer, equivalent to block for heap
    # with blocks_per_bucket=1)
    #

    @property
    def blocks_per_bucket(self):
        return self._blocks_per_bucket

    def bucket_count_at_level(self, l):
        return calculate_bucket_count_in_heap_at_level(self.k, l)
    def first_bucket_at_level(self, l):
        if l > 0:
            return calculate_bucket_count_in_heap_with_height(self.k, l-1)
        return 0
    def last_bucket_at_level(self, l):
        return calculate_bucket_count_in_heap_with_height(self.k, l) - 1
    def random_bucket_up_to_level(self, l):
        return self.random.randint(self.first_bucket_at_level(0),
                                   self.last_bucket_at_level(l))
    def random_bucket_at_level(self, l):
        return self.random.randint(self.first_bucket_at_level(l),
                                   self.first_bucket_at_level(l+1)-1)

    #
    # Nodes (a class that helps with heap path calculations)
    #

    def root_node(self):
        return self.first_node_at_level(0)
    def node_count_at_level(self, l):
        return self.bucket_count_at_level(l)
    def first_node_at_level(self, l):
        return self.Node(self.first_bucket_at_level(l))
    def last_node_at_level(self, l):
        return self.Node(self.last_bucket_at_level(l))
    def random_node_up_to_level(self, l):
        return self.Node(self.random_bucket_up_to_level(l))
    def random_node_at_level(self, l):
        return self.Node(self.random_bucket_at_level(l))

    #
    # Block (0-based integer)
    #

    def bucket_to_block(self, b):
        assert b >= 0
        return b * self.blocks_per_bucket
    def block_to_bucket(self, s):
        assert s >= 0
        return s//self.blocks_per_bucket
    def first_block_in_bucket(self, b):
        return self.bucket_to_block(b)
    def last_block_in_bucket(self, b):
        return self.bucket_to_block(b) + self.blocks_per_bucket - 1
    def block_count_at_level(self, l):
        return self.bucket_count_at_level(l) * self.blocks_per_bucket
    def first_block_at_level(self, l):
        return self.bucket_to_block(self.first_bucket_at_level(l))
    def last_block_at_level(self, l):
        return self.bucket_to_block(self.first_bucket_at_level(l+1)) - 1

class SizedVirtualHeap(VirtualHeap):

    def __init__(self, k, height, blocks_per_bucket=1):
        super(SizedVirtualHeap, self).\
            __init__(k, blocks_per_bucket=blocks_per_bucket)
        self._height = height

    #
    # Size properties
    #

    @property
    def levels(self):
        return self._height + 1
    @property
    def height(self):
        return self._height

    #
    # Buckets (0-based integer, equivalent to block for heap
    # with blocks_per_bucket=1)
    #

    def bucket_count(self):
        return calculate_bucket_count_in_heap_with_height(self.k,
                                                          self.height)
    def leaf_bucket_count(self):
        return calculate_leaf_bucket_count_in_heap_with_height(self.k,
                                                               self.height)
    def first_leaf_bucket(self):
        return self.first_bucket_at_level(self.height)
    def last_leaf_bucket(self):
        return self.last_bucket_at_level(self.height)
    def random_bucket(self):
        return self.random.randint(self.first_bucket_at_level(0),
                                   self.last_leaf_bucket())
    def random_leaf_bucket(self):
        return self.random_bucket_at_level(self.height)

    #
    # Nodes (a class that helps with heap path calculations)
    #

    def is_nil_node(self, n):
        return n.bucket >= self.bucket_count()
    def node_count(self):
        return self.bucket_count()
    def leaf_node_count(self):
        return self.leaf_bucket_count()
    def first_leaf_node(self):
        return self.Node(self.first_leaf_bucket())
    def last_leaf_node(self):
        return self.Node(self.last_leaf_bucket())
    def random_leaf_node(self):
        return self.Node(self.random_leaf_bucket())
    def random_node(self):
        return self.Node(self.random_bucket())

    #
    # Block (0-based integer)
    #

    def block_count(self):
        return self.bucket_count() * self.blocks_per_bucket
    def leaf_block_count(self):
        return self.leaf_bucket_count() * self.blocks_per_bucket
    def first_leaf_block(self):
        return self.first_block_in_bucket(self.first_leaf_bucket())
    def last_leaf_block(self):
        return self.last_block_in_bucket(self.last_leaf_bucket())

    #
    # Visualization
    #

    def write_as_dot(self, f, data=None, max_levels=None):
        "Write the tree in the dot language format to f."
        assert (max_levels is None) or (max_levels >= 0)
        def visit_node(n, levels):
            lbl = "{"
            if data is None:
                if self.k <= max_k_labeled:
                    lbl = repr(n.label()).\
                          replace("{","\{").\
                          replace("}","\}").\
                          replace("|","\|").\
                          replace("<","\<").\
                          replace(">","\>")
                else:
                    lbl = str(n)
            else:
                s = self.bucket_to_block(n.bucket)
                for i in range(self.blocks_per_bucket):
                    lbl += "{%s}" % (data[s+i])
                    if i + 1 != self.blocks_per_bucket:
                        lbl += "|"
            lbl += "}"
            f.write("  %s [penwidth=%s,label=\"%s\"];\n"
                    % (n.bucket, 1, lbl))
            levels += 1
            if (max_levels is None) or (levels <= max_levels):
                for i in range(self.k):
                    cn = n.child_node(i)
                    if not self.is_nil_node(cn):
                        visit_node(cn, levels)
                        f.write("  %s -> %s ;\n" % (n.bucket, cn.bucket))

        f.write("// Created by SizedVirtualHeap.write_as_dot(...)\n")
        f.write("digraph heaptree {\n")
        f.write("node [shape=record]\n")

        if (max_levels is None) or (max_levels > 0):
            visit_node(self.root_node(), 1)
        f.write("}\n")

    def save_image_as_pdf(self, filename, data=None, max_levels=None):
        "Write the heap as PDF file."
        assert (max_levels is None) or (max_levels >= 0)
        import os
        if not filename.endswith('.pdf'):
            filename = filename+'.pdf'
        tmpfd, tmpname = tempfile.mkstemp(suffix='dot')
        with open(tmpname, 'w') as f:
            self.write_as_dot(f, data=data, max_levels=max_levels)
        os.close(tmpfd)
        try:
            subprocess.call(['dot',
                             tmpname,
                             '-Tpdf',
                             '-o',
                             ('%s'%filename)])
        except OSError:
            sys.stderr.write(
                "DOT -> PDF conversion failed. See DOT file: %s\n"
                % (tmpname))
            return False
        os.remove(tmpname)
        return True
<commit_msg>adding a few more helper properties to virtual heap<commit_after>__all__ = ("VirtualHeap",
           "SizedVirtualHeap")

import os
import sys
import subprocess
import random
import string
import tempfile

from six.moves import range

from pyoram.util._virtual_heap_helper import lib as _clib
from pyoram.util.misc import log2floor

numerals = ''.join([c for c in string.printable \
                  if ((c not in string.whitespace) and \
                      (c != '+') and (c != '-') and \
                      (c != '"') and (c != "'") and \
                      (c != '\\') and (c != '/'))])
numeral_index = dict((c,i) for i,c in enumerate(numerals))

# The maximum heap base for which base k labels
# can be produced.
max_k_labeled = len(numerals)

def base10_integer_to_basek_string(k, x):
    """Convert an integer into a base k string."""
    if not (2 <= k <= max_k_labeled):
        raise ValueError("k must be in range [2, %d]: %s"
                         % (max_k_labeled, k))
    return ((x == 0) and numerals[0]) or \
        (base10_integer_to_basek_string(k, x // k).\
         lstrip(numerals[0]) + numerals[x % k])

def basek_string_to_base10_integer(k, x):
    """Convert a base k string into an integer."""
    assert 1 < k <= max_k_labeled
    return sum(numeral_index[c]*(k**i)
               for i, c in enumerate(reversed(x)))

# _clib defines a faster version of this function
def calculate_bucket_level(k, b):
    """
    Calculate the level in which a 0-based bucket
    lives inside of a k-ary heap.
    """
    assert k >= 2
    if k == 2:
        return log2floor(b+1)
    v = (k - 1) * (b + 1) + 1
    h = 0
    while k**(h+1) < v:
        h += 1
    return h

# _clib defines a faster version of this function
def calculate_last_common_level(k, b1, b2):
    """
    Calculate the highest level after which the
    paths from the root to these buckets diverge.
    """
    l1 = calculate_bucket_level(k, b1)
    l2 = calculate_bucket_level(k, b2)
    while l1 > l2:
        b1 = (b1-1)//k
        l1 -= 1
    while l2 > l1:
        b2 = (b2-1)//k
        l2 -= 1
    while b1 != b2:
        b1 = (b1-1)//k
        b2 = (b2-1)//k
        l1 -= 1
    return l1

def calculate_necessary_heap_height(k, n):
    """
    Calculate the necessary k-ary heap height
    to store n buckets.
    """
    assert n >= 1
    return calculate_bucket_level(k, n-1)

def calculate_bucket_count_in_heap_with_height(k, h):
    """
    Calculate the number of buckets in a
    k-ary heap of height h.
    """
    assert h >= 0
    return ((k**(h+1)) - 1) // (k - 1)

def calculate_bucket_count_in_heap_at_level(k, l):
    """
    Calculate the number of buckets in a
    k-ary heap at level l.
    """
    assert l >= 0
    return k**l

def calculate_leaf_bucket_count_in_heap_with_height(k, h):
    """
    Calculate the number of buckets in the
    leaf-level of a k-ary heap of height h.
    """
    return calculate_bucket_count_in_heap_at_level(k, h)

def create_node_type(k):

    class VirtualHeapNode(object):
        __slots__ = ("bucket", "level")
        def __init__(self, bucket):
            assert bucket >= 0
            self.bucket = bucket
            self.level = _clib.calculate_bucket_level(self.k, self.bucket)

        def __hash__(self):
            return self.bucket.__hash__()
        def __lt__(self, other):
            return self.bucket < other
        def __le__(self, other):
            return self.bucket <= other
        def __eq__(self, other):
            return self.bucket == other
        def __ne__(self, other):
            return self.bucket != other
        def __gt__(self, other):
            return self.bucket > other
        def __ge__(self, other):
            return self.bucket >= other
        def last_common_level(self, n):
            return _clib.calculate_last_common_level(self.k,
                                                     self.bucket,
                                                     n.bucket)
        def child_node(self, c):
            assert type(c) is int
            assert 0 <= c < self.k
            return VirtualHeapNode(self.k * self.bucket + 1 + c)
        def parent_node(self):
            if self.bucket != 0:
                return VirtualHeapNode((self.bucket - 1)//self.k)
            return None
        def ancestor_node_at_level(self, level):
            if level > self.level:
                return None
            current = self
            while current.level != level:
                current = current.parent_node()
            return current
        def bucket_path_to_root(self):
            bucket = self.bucket
            yield bucket
            while bucket != 0:
                bucket = (bucket - 1)//self.k
                yield bucket
        def bucket_path_from_root(self):
            return list(reversed(list(self.bucket_path_to_root())))

        #
        # Expensive Functions
        #
        def __repr__(self):
            try:
                label = self.label()
            except ValueError:
                # presumably, k is too large
                label = "<unknown>"
            return ("VirtualHeapNode(k=%s, bucket=%s, level=%s, label=%r)"
                    % (self.k, self.bucket, self.level, label))
        def __str__(self):
            """Returns a tuple (<level>, <bucket offset within level>)."""
            if self.bucket != 0:
                return ("(%s, %s)"
                        % (self.level,
                           self.bucket -
                           calculate_bucket_count_in_heap_with_height(self.k,
                                                                self.level-1)))
            assert self.level == 0
            return "(0, 0)"

        def label(self):
            assert 0 <= self.bucket
            if self.level == 0:
                return ''
            b_offset = self.bucket - \
                       calculate_bucket_count_in_heap_with_height(self.k,
                                                            self.level-1)
            basek = base10_integer_to_basek_string(self.k, b_offset)
            return basek.zfill(self.level)

        def is_node_on_path(self, n):
            if n.level <= self.level:
                n_label = n.label()
                if n_label == "":
                    return True
                return self.label().startswith(n_label)
            return False

    VirtualHeapNode.k = k

    return VirtualHeapNode

class VirtualHeap(object):

    clib = _clib
    random = random.SystemRandom()

    def __init__(self, k, blocks_per_bucket=1):
        assert 1 < k
        assert blocks_per_bucket >= 1
        self._k = k
        self._blocks_per_bucket = blocks_per_bucket
        self.Node = create_node_type(k)

    @property
    def k(self):
        return self._k

    def node_label_to_bucket(self, label):
        if len(label) > 0:
            return \
                (calculate_bucket_count_in_heap_with_height(self.k,
                                                      len(label)-1) +
                 basek_string_to_base10_integer(self.k, label))
        return 0

    #
    # Buckets (0-based integer, equivalent to block for heap
    # with blocks_per_bucket=1)
    #

    @property
    def blocks_per_bucket(self):
        return self._blocks_per_bucket

    def bucket_count_at_level(self, l):
        return calculate_bucket_count_in_heap_at_level(self.k, l)
    def first_bucket_at_level(self, l):
        if l > 0:
            return calculate_bucket_count_in_heap_with_height(self.k, l-1)
        return 0
    def last_bucket_at_level(self, l):
        return calculate_bucket_count_in_heap_with_height(self.k, l) - 1
    def random_bucket_up_to_level(self, l):
        return self.random.randint(self.first_bucket_at_level(0),
                                   self.last_bucket_at_level(l))
    def random_bucket_at_level(self, l):
        return self.random.randint(self.first_bucket_at_level(l),
                                   self.first_bucket_at_level(l+1)-1)

    #
    # Nodes (a class that helps with heap path calculations)
    #

    def root_node(self):
        return self.first_node_at_level(0)
    def node_count_at_level(self, l):
        return self.bucket_count_at_level(l)
    def first_node_at_level(self, l):
        return self.Node(self.first_bucket_at_level(l))
    def last_node_at_level(self, l):
        return self.Node(self.last_bucket_at_level(l))
    def random_node_up_to_level(self, l):
        return self.Node(self.random_bucket_up_to_level(l))
    def random_node_at_level(self, l):
        return self.Node(self.random_bucket_at_level(l))

    #
    # Block (0-based integer)
    #

    def bucket_to_block(self, b):
        assert b >= 0
        return b * self.blocks_per_bucket
    def block_to_bucket(self, s):
        assert s >= 0
        return s//self.blocks_per_bucket
    def first_block_in_bucket(self, b):
        return self.bucket_to_block(b)
    def last_block_in_bucket(self, b):
        return self.bucket_to_block(b) + self.blocks_per_bucket - 1
    def block_count_at_level(self, l):
        return self.bucket_count_at_level(l) * self.blocks_per_bucket
    def first_block_at_level(self, l):
        return self.bucket_to_block(self.first_bucket_at_level(l))
    def last_block_at_level(self, l):
        return self.bucket_to_block(self.first_bucket_at_level(l+1)) - 1

class SizedVirtualHeap(VirtualHeap):

    def __init__(self, k, height, blocks_per_bucket=1):
        super(SizedVirtualHeap, self).\
            __init__(k, blocks_per_bucket=blocks_per_bucket)
        self._height = height

    #
    # Size properties
    #
    @property
    def height(self):
        return self._height
    @property
    def levels(self):
        return self.height + 1
    @property
    def first_level(self):
        return 0
    @property
    def last_level(self):
        return self.height

    #
    # Buckets (0-based integer, equivalent to block for heap
    # with blocks_per_bucket=1)
    #

    def bucket_count(self):
        return calculate_bucket_count_in_heap_with_height(self.k,
                                                          self.height)
    def leaf_bucket_count(self):
        return calculate_leaf_bucket_count_in_heap_with_height(self.k,
                                                               self.height)
    def first_leaf_bucket(self):
        return self.first_bucket_at_level(self.height)
    def last_leaf_bucket(self):
        return self.last_bucket_at_level(self.height)
    def random_bucket(self):
        return self.random.randint(self.first_bucket_at_level(0),
                                   self.last_leaf_bucket())
    def random_leaf_bucket(self):
        return self.random_bucket_at_level(self.height)

    #
    # Nodes (a class that helps with heap path calculations)
    #

    def is_nil_node(self, n):
        return n.bucket >= self.bucket_count()
    def node_count(self):
        return self.bucket_count()
    def leaf_node_count(self):
        return self.leaf_bucket_count()
    def first_leaf_node(self):
        return self.Node(self.first_leaf_bucket())
    def last_leaf_node(self):
        return self.Node(self.last_leaf_bucket())
    def random_leaf_node(self):
        return self.Node(self.random_leaf_bucket())
    def random_node(self):
        return self.Node(self.random_bucket())

    #
    # Block (0-based integer)
    #

    def block_count(self):
        return self.bucket_count() * self.blocks_per_bucket
    def leaf_block_count(self):
        return self.leaf_bucket_count() * self.blocks_per_bucket
    def first_leaf_block(self):
        return self.first_block_in_bucket(self.first_leaf_bucket())
    def last_leaf_block(self):
        return self.last_block_in_bucket(self.last_leaf_bucket())

    #
    # Visualization
    #

    def write_as_dot(self, f, data=None, max_levels=None):
        "Write the tree in the dot language format to f."
        assert (max_levels is None) or (max_levels >= 0)
        def visit_node(n, levels):
            lbl = "{"
            if data is None:
                if self.k <= max_k_labeled:
                    lbl = repr(n.label()).\
                          replace("{","\{").\
                          replace("}","\}").\
                          replace("|","\|").\
                          replace("<","\<").\
                          replace(">","\>")
                else:
                    lbl = str(n)
            else:
                s = self.bucket_to_block(n.bucket)
                for i in range(self.blocks_per_bucket):
                    lbl += "{%s}" % (data[s+i])
                    if i + 1 != self.blocks_per_bucket:
                        lbl += "|"
            lbl += "}"
            f.write("  %s [penwidth=%s,label=\"%s\"];\n"
                    % (n.bucket, 1, lbl))
            levels += 1
            if (max_levels is None) or (levels <= max_levels):
                for i in range(self.k):
                    cn = n.child_node(i)
                    if not self.is_nil_node(cn):
                        visit_node(cn, levels)
                        f.write("  %s -> %s ;\n" % (n.bucket, cn.bucket))

        f.write("// Created by SizedVirtualHeap.write_as_dot(...)\n")
        f.write("digraph heaptree {\n")
        f.write("node [shape=record]\n")

        if (max_levels is None) or (max_levels > 0):
            visit_node(self.root_node(), 1)
        f.write("}\n")

    def save_image_as_pdf(self, filename, data=None, max_levels=None):
        "Write the heap as PDF file."
        assert (max_levels is None) or (max_levels >= 0)
        import os
        if not filename.endswith('.pdf'):
            filename = filename+'.pdf'
        tmpfd, tmpname = tempfile.mkstemp(suffix='dot')
        with open(tmpname, 'w') as f:
            self.write_as_dot(f, data=data, max_levels=max_levels)
        os.close(tmpfd)
        try:
            subprocess.call(['dot',
                             tmpname,
                             '-Tpdf',
                             '-o',
                             ('%s'%filename)])
        except OSError:
            sys.stderr.write(
                "DOT -> PDF conversion failed. See DOT file: %s\n"
                % (tmpname))
            return False
        os.remove(tmpname)
        return True
<|endoftext|>


<commit_before>import sys
from http import cookies

# Cookie pickling bug is fixed in Python 2.7.9 and Python 3.4.3+
# http://bugs.python.org/issue22775
cookie_pickles_properly = (
    (sys.version_info[:2] == (2, 7) and sys.version_info >= (2, 7, 9)) or
    sys.version_info >= (3, 4, 3)
)

if cookie_pickles_properly:
    SimpleCookie = cookies.SimpleCookie
else:
    Morsel = cookies.Morsel

    class SimpleCookie(cookies.SimpleCookie):
        if not cookie_pickles_properly:
            def __setitem__(self, key, value):
                # Apply the fix from http://bugs.python.org/issue22775 where
                # it's not fixed in Python itself
                if isinstance(value, Morsel):
                    # allow assignment of constructed Morsels (e.g. for pickling)
                    dict.__setitem__(self, key, value)
                else:
                    super(SimpleCookie, self).__setitem__(key, value)

def parse_cookie(cookie):
    """
    Return a dictionary parsed from a `Cookie:` header string.
    """
    cookiedict = {}
    for chunk in cookie.split(str(';')):
        if str('=') in chunk:
            key, val = chunk.split(str('='), 1)
        else:
            # Assume an empty name per
            # https://bugzilla.mozilla.org/show_bug.cgi?id=169091
            key, val = str(''), chunk
        key, val = key.strip(), val.strip()
        if key or val:
            # unquote using Python's algorithm.
            cookiedict[key] = cookies._unquote(val)
    return cookiedict
<commit_msg>Refs #23919 -- Removed Python 2 version check in django.http.cookie.<commit_after>import sys
from http import cookies

# Cookie pickling bug is fixed in Python 3.4.3+
# http://bugs.python.org/issue22775
if sys.version_info >= (3, 4, 3):
    SimpleCookie = cookies.SimpleCookie
else:
    Morsel = cookies.Morsel

    class SimpleCookie(cookies.SimpleCookie):
        def __setitem__(self, key, value):
            if isinstance(value, Morsel):
                # allow assignment of constructed Morsels (e.g. for pickling)
                dict.__setitem__(self, key, value)
            else:
                super(SimpleCookie, self).__setitem__(key, value)

def parse_cookie(cookie):
    """
    Return a dictionary parsed from a `Cookie:` header string.
    """
    cookiedict = {}
    for chunk in cookie.split(str(';')):
        if str('=') in chunk:
            key, val = chunk.split(str('='), 1)
        else:
            # Assume an empty name per
            # https://bugzilla.mozilla.org/show_bug.cgi?id=169091
            key, val = str(''), chunk
        key, val = key.strip(), val.strip()
        if key or val:
            # unquote using Python's algorithm.
            cookiedict[key] = cookies._unquote(val)
    return cookiedict
<|endoftext|>