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

import io
import collections
import dataclasses
import os

import flax
import jax
import jax.numpy as jnp
import numpy as np
import scipy
from tensorflow.io import gfile  # pylint: disable=import-error

def _flatten_dict(d, parent_key='', sep='/'):
  """Flattens a dictionary, keeping empty leaves."""
  items = []
  for k, v in d.items():
    path = parent_key + sep + k if parent_key else k
    if isinstance(v, collections.MutableMapping):
      items.extend(_flatten_dict(v, path, sep=sep).items())
    else:
      items.append((path, v))

  # Keeps the empty dict if it was set explicitly.
  if parent_key and not d:
    items.append((parent_key, {}))

  return dict(items)

def inspect_params(*,
                   params,
                   expected,
                   logger,
                   fail_if_extra=True,
                   fail_if_missing=True):
  """Inspects whether the params are consistent with the expected keys."""
  params_flat = _flatten_dict(params)
  expected_flat = _flatten_dict(expected)
  missing_keys = expected_flat.keys() - params_flat.keys()
  extra_keys = params_flat.keys() - expected_flat.keys()

  # Adds back empty dict explicitly, to support layers without weights.
  # Context: FLAX ignores empty dict during serialization.
  empty_keys = set()
  for k in missing_keys:
    if isinstance(expected_flat[k], dict) and not expected_flat[k]:
      params[k] = {}
      empty_keys.add(k)
  missing_keys -= empty_keys

  if empty_keys:
    logger.warning('Inspect recovered empty keys:\n%s', empty_keys)
  if missing_keys:
    logger.info('Inspect missing keys:\n%s', missing_keys)
  if extra_keys:
    logger.info('Inspect extra keys:\n%s', extra_keys)

  if (missing_keys and fail_if_missing) or (extra_keys and fail_if_extra):
    raise ValueError(f'Missing params from checkpoint: {missing_keys}.\n'
                     f'Extra params in checkpoint: {extra_keys}.\n'
                     f'Restored params from checkpoint: {params_flat.keys()}.\n'
                     f'Expected params from code: {expected_flat.keys()}.')
  return params

def recover_tree(keys, values):
  """Recovers a tree as a nested dict from flat names and values.

  This function is useful to analyze checkpoints that are without need to access
  the exact source code of the experiment. In particular, it can be used to
  extract an reuse various subtrees of the scheckpoint, e.g. subtree of
  parameters.

  Args:
    keys: a list of keys, where '/' is used as separator between nodes.
    values: a list of leaf values.

  Returns:
    A nested tree-like dict.
  """
  tree = {}
  sub_trees = collections.defaultdict(list)
  for k, v in zip(keys, values):
    if '/' not in k:
      tree[k] = v
    else:
      k_left, k_right = k.split('/', 1)
      sub_trees[k_left].append((k_right, v))
  for k, kv_pairs in sub_trees.items():
    k_subtree, v_subtree = zip(*kv_pairs)
    tree[k] = recover_tree(k_subtree, v_subtree)
  return tree

def _traverse_with_names(tree):
  """Traverses nested dicts/dataclasses and emits (leaf_name, leaf_val)."""
  if dataclasses.is_dataclass(tree):
    tree = flax.serialization.to_state_dict(tree)
  if isinstance(tree, dict):
    keys = sorted(tree.keys())
    for key in keys:
      for path, v in _traverse_with_names(tree[key]):
        yield (key + '/' + path).rstrip('/'), v
  else:
    yield '', tree

def tree_flatten_with_names(tree):
  """Populates tree_flatten with leaf names.

  This function populates output of tree_flatten with leaf names, using a
  custom traversal that produces names is provided. The custom traversal does
  NOT have to traverse tree in the same order as jax, as we take care of
  automatically aligning jax' and custom traversals.

  Args:
    tree: python tree.

  Returns:
    A list of values with names: [(name, value), ...]
  """
  vals, tree_def = jax.tree_flatten(tree)

  # "Fake" token tree that is use to track jax internal tree traversal and
  # adjust our custom tree traversal to be compatible with it.
  tokens = range(len(vals))
  token_tree = tree_def.unflatten(tokens)
  val_names, perm = zip(*_traverse_with_names(token_tree))
  inv_perm = np.argsort(perm)

  # Custom traverasal should visit the same number of leaves.
  assert len(val_names) == len(vals)

  return [(val_names[i], v) for i, v in zip(inv_perm, vals)], tree_def

def save(data, path):
  """Util for checkpointing: saves jax pytree objects to the disk.

  These checkpoints can later be recovered with `load()`.

  Args:
    data: arbitrary jax pytree to be saved.
    path: a path to save the data.
  """
  names_and_vals, _ = tree_flatten_with_names(data)
  io_buffer = io.BytesIO()

  # savez uses `seek()` API call, which is not supported by cns. Thus, we first
  # write the checkpoint to the temp buffer and then write it to the disk.
  np.savez(io_buffer, **{k: v for k, v in names_and_vals})

  # In order to be robust to interruptions we first save checkpoint to the
  # temporal file and then move to actual path name.
  path_tmp = path + '-TEMPORARY'
  gfile.makedirs(os.path.dirname(path))
  with gfile.GFile(path_tmp, 'wb') as f:
    f.write(io_buffer.getvalue())
  gfile.rename(path_tmp, path, overwrite=True)

def load(path):
  """Loads params from a checkpoint previously stored with `save()`."""
  with gfile.GFile(path, 'rb') as f:
    ckpt_dict = np.load(f, allow_pickle=False)
    keys, values = zip(*list(ckpt_dict.items()))
  return recover_tree(keys, values)

def load_pretrained(*, pretrained_path, init_params, model_config, logger):
  """Loads/converts a pretrained checkpoint for fine tuning.

  Args:
    logger: Logger to use to output diagnostic messages.
    init_params: Parameters from model. Will be used for the head of the model
      and to verify that the model is compatible with the stored checkpoint.
    init_file: File pointing to pretrained checkpoint.
    model_config: Configuration of the model. Will be used to configure the
      head and rescale the position embeddings.

  Returns:
    Parameters like `init_params`, but loaded with pretrained weights from
    `init_file` and adapted accordingly.
  """

  restored_params = inspect_params(
      params=load(pretrained_path),
      expected=init_params,
      logger=logger,
      fail_if_extra=False,
      fail_if_missing=False)

  # The following allows implementing fine-tuning head variants depending on the
  # value of `representation_size` in the fine-tuning job:
  # - `None` : drop the whole head and attach a nn.Linear.
  # - same number as in pre-training means : keep the head but reset the last
  #    layer (logits) for the new task.
  if model_config.representation_size is None:
    if 'pre_logits' in restored_params:
      logger.info('load_pretrained: drop-head variant')
      restored_params['pre_logits'] = {}
  restored_params['head']['kernel'] = init_params['head']['kernel']
  restored_params['head']['bias'] = init_params['head']['bias']

  if 'posembed_input' in restored_params.get('Transformer', {}):
    # Rescale the grid of position embeddings. Param shape is (1,N,1024)
    posemb = restored_params['Transformer']['posembed_input']['pos_embedding']
    posemb_new = init_params['Transformer']['posembed_input']['pos_embedding']
    if posemb.shape != posemb_new.shape:
      logger.info('load_pretrained: resized variant: %s to %s', posemb.shape,
                  posemb_new.shape)
      ntok_new = posemb_new.shape[1]

      if model_config.classifier == 'token':
        posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
        ntok_new -= 1
      else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]

      gs_old = int(np.sqrt(len(posemb_grid)))
      gs_new = int(np.sqrt(ntok_new))
      logger.info('load_pretrained: grid-size from %s to %s', gs_old, gs_new)
      posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)

      zoom = (gs_new / gs_old, gs_new / gs_old, 1)
      posemb_grid = scipy.ndimage.zoom(posemb_grid, zoom, order=1)
      posemb_grid = posemb_grid.reshape(1, gs_new * gs_new, -1)
      posemb = jnp.array(np.concatenate([posemb_tok, posemb_grid], axis=1))
      restored_params['Transformer']['posembed_input']['pos_embedding'] = posemb

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

import io
import collections
import dataclasses
import os

import flax
import jax
import jax.numpy as jnp
import numpy as np
import scipy
from tensorflow.io import gfile  # pylint: disable=import-error

def _flatten_dict(d, parent_key='', sep='/'):
  """Flattens a dictionary, keeping empty leaves."""
  items = []
  for k, v in d.items():
    path = parent_key + sep + k if parent_key else k
    if isinstance(v, collections.MutableMapping):
      items.extend(_flatten_dict(v, path, sep=sep).items())
    else:
      items.append((path, v))

  # Keeps the empty dict if it was set explicitly.
  if parent_key and not d:
    items.append((parent_key, {}))

  return dict(items)

def inspect_params(*,
                   params,
                   expected,
                   logger,
                   fail_if_extra=True,
                   fail_if_missing=True):
  """Inspects whether the params are consistent with the expected keys."""
  params_flat = _flatten_dict(params)
  expected_flat = _flatten_dict(expected)
  missing_keys = expected_flat.keys() - params_flat.keys()
  extra_keys = params_flat.keys() - expected_flat.keys()

  # Adds back empty dict explicitly, to support layers without weights.
  # Context: FLAX ignores empty dict during serialization.
  empty_keys = set()
  for k in missing_keys:
    if isinstance(expected_flat[k], dict) and not expected_flat[k]:
      params[k] = {}
      empty_keys.add(k)
  missing_keys -= empty_keys

  if empty_keys:
    logger.warning('Inspect recovered empty keys:\n%s', empty_keys)
  if missing_keys:
    logger.info('Inspect missing keys:\n%s', missing_keys)
  if extra_keys:
    logger.info('Inspect extra keys:\n%s', extra_keys)

  if (missing_keys and fail_if_missing) or (extra_keys and fail_if_extra):
    raise ValueError(f'Missing params from checkpoint: {missing_keys}.\n'
                     f'Extra params in checkpoint: {extra_keys}.\n'
                     f'Restored params from checkpoint: {params_flat.keys()}.\n'
                     f'Expected params from code: {expected_flat.keys()}.')
  return params

def recover_tree(keys, values):
  """Recovers a tree as a nested dict from flat names and values.

  This function is useful to analyze checkpoints that are without need to access
  the exact source code of the experiment. In particular, it can be used to
  extract an reuse various subtrees of the scheckpoint, e.g. subtree of
  parameters.

  Args:
    keys: a list of keys, where '/' is used as separator between nodes.
    values: a list of leaf values.

  Returns:
    A nested tree-like dict.
  """
  tree = {}
  sub_trees = collections.defaultdict(list)
  for k, v in zip(keys, values):
    if '/' not in k:
      tree[k] = v
    else:
      k_left, k_right = k.split('/', 1)
      sub_trees[k_left].append((k_right, v))
  for k, kv_pairs in sub_trees.items():
    k_subtree, v_subtree = zip(*kv_pairs)
    tree[k] = recover_tree(k_subtree, v_subtree)
  return tree

def _traverse_with_names(tree):
  """Traverses nested dicts/dataclasses and emits (leaf_name, leaf_val)."""
  if dataclasses.is_dataclass(tree):
    tree = flax.serialization.to_state_dict(tree)
  if isinstance(tree, dict):
    keys = sorted(tree.keys())
    for key in keys:
      for path, v in _traverse_with_names(tree[key]):
        yield (key + '/' + path).rstrip('/'), v
  else:
    yield '', tree

def tree_flatten_with_names(tree):
  """Populates tree_flatten with leaf names.

  This function populates output of tree_flatten with leaf names, using a
  custom traversal that produces names is provided. The custom traversal does
  NOT have to traverse tree in the same order as jax, as we take care of
  automatically aligning jax' and custom traversals.

  Args:
    tree: python tree.

  Returns:
    A list of values with names: [(name, value), ...]
  """
  vals, tree_def = jax.tree_flatten(tree)

  # "Fake" token tree that is use to track jax internal tree traversal and
  # adjust our custom tree traversal to be compatible with it.
  tokens = range(len(vals))
  token_tree = tree_def.unflatten(tokens)
  val_names, perm = zip(*_traverse_with_names(token_tree))
  inv_perm = np.argsort(perm)

  # Custom traverasal should visit the same number of leaves.
  assert len(val_names) == len(vals)

  return [(val_names[i], v) for i, v in zip(inv_perm, vals)], tree_def

def save(data, path):
  """Util for checkpointing: saves jax pytree objects to the disk.

  These checkpoints can later be recovered with `load()`.

  Args:
    data: arbitrary jax pytree to be saved.
    path: a path to save the data.
  """
  names_and_vals, _ = tree_flatten_with_names(data)
  io_buffer = io.BytesIO()

  # savez uses `seek()` API call, which is not supported by cns. Thus, we first
  # write the checkpoint to the temp buffer and then write it to the disk.
  np.savez(io_buffer, **{k: v for k, v in names_and_vals})

  # In order to be robust to interruptions we first save checkpoint to the
  # temporal file and then move to actual path name.
  path_tmp = path + '-TEMPORARY'
  gfile.makedirs(os.path.dirname(path))
  with gfile.GFile(path_tmp, 'wb') as f:
    f.write(io_buffer.getvalue())
  gfile.rename(path_tmp, path, overwrite=True)

def load(path):
  """Loads params from a checkpoint previously stored with `save()`."""
  with gfile.GFile(path, 'rb') as f:
    ckpt_dict = np.load(f, allow_pickle=False)
    keys, values = zip(*list(ckpt_dict.items()))
  return recover_tree(keys, values)

def load_pretrained(*, pretrained_path, init_params, model_config, logger):
  """Loads/converts a pretrained checkpoint for fine tuning.

  Args:
    logger: Logger to use to output diagnostic messages.
    init_params: Parameters from model. Will be used for the head of the model
      and to verify that the model is compatible with the stored checkpoint.
    pretrained_path: File pointing to pretrained checkpoint.
    model_config: Configuration of the model. Will be used to configure the
      head and rescale the position embeddings.

  Returns:
    Parameters like `init_params`, but loaded with pretrained weights from
    `pretrained_path` and adapted accordingly.
  """

  restored_params = inspect_params(
      params=load(pretrained_path),
      expected=init_params,
      logger=logger,
      fail_if_extra=False,
      fail_if_missing=False)

  # The following allows implementing fine-tuning head variants depending on the
  # value of `representation_size` in the fine-tuning job:
  # - `None` : drop the whole head and attach a nn.Linear.
  # - same number as in pre-training means : keep the head but reset the last
  #    layer (logits) for the new task.
  if model_config.representation_size is None:
    if 'pre_logits' in restored_params:
      logger.info('load_pretrained: drop-head variant')
      restored_params['pre_logits'] = {}
  restored_params['head']['kernel'] = init_params['head']['kernel']
  restored_params['head']['bias'] = init_params['head']['bias']

  if 'posembed_input' in restored_params.get('Transformer', {}):
    # Rescale the grid of position embeddings. Param shape is (1,N,1024)
    posemb = restored_params['Transformer']['posembed_input']['pos_embedding']
    posemb_new = init_params['Transformer']['posembed_input']['pos_embedding']
    if posemb.shape != posemb_new.shape:
      logger.info('load_pretrained: resized variant: %s to %s', posemb.shape,
                  posemb_new.shape)
      ntok_new = posemb_new.shape[1]

      if model_config.classifier == 'token':
        posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
        ntok_new -= 1
      else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]

      gs_old = int(np.sqrt(len(posemb_grid)))
      gs_new = int(np.sqrt(ntok_new))
      logger.info('load_pretrained: grid-size from %s to %s', gs_old, gs_new)
      posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)

      zoom = (gs_new / gs_old, gs_new / gs_old, 1)
      posemb_grid = scipy.ndimage.zoom(posemb_grid, zoom, order=1)
      posemb_grid = posemb_grid.reshape(1, gs_new * gs_new, -1)
      posemb = jnp.array(np.concatenate([posemb_tok, posemb_grid], axis=1))
      restored_params['Transformer']['posembed_input']['pos_embedding'] = posemb

  return restored_params
<|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8
from __future__ import unicode_literals, absolute_import, print_function
from flask_restaction import Resource
from jinja2 import Template
import codecs
from pkg_resources import resource_string
import os
from flask import json

def parse_api(api):
    """parse api"""
    data = {"url_prefix": "", "auth_header": ""}
    data["docs"] = api.docs
    if api.url_prefix:
        data["url_prefix"] = api.url_prefix
    if api.auth:
        data["auth_header"] = api.auth.auth_header
    resources = []
    for name, res in api.resources.items():
        reaource = {"name": name, "docs": res["docs"]}
        actions = [dict(x.__dict__) for x in res["actions"]]
        reaource["actions"] = actions
        resources.append(reaource)
    data["resources"] = resources
    return data

class ApiInfo(Resource):
    """ApiInfo

    :param api: Api
    """
    action = {
        "action": "str",  # post_login
        "method": "str",  # post
        "url": "str",  # /api/user/login
        "endpoint": "str",  # user@login
        "docs": "str",  # docs of post_login
        "inputs": "str",  # schema_inputs[action]
        "outputs": "str",  # schema_outputs[action]
    }
    schema_outputs = {
        "get": {
            "url_prefix": "str",  # /api
            "auth_header": "str",  # Authorization
            "docs": "str",  # docs of api
            "resources": [{
                "name": "str",  # user
                "docs": "str",  # docs of user
                "actions": [action]
            }],
        }
    }

    def __init__(self, api):
        self.data = parse_api(api)

    def get(self):
        return self.data

class Gen(object):
    """Generate tools for docs res.js

    if dest param of generate methods is not absolute path,
    then join app.root_path and dest as dest path.
    :param api: Api
    """

    def __init__(self, api):
        self.api = api
        self.data = parse_api(api)
        app = self.api.app
        if not os.path.exists(app.static_folder):
            os.makedirs(app.static_folder)

    def _read_file(self, *paths):
        strs = [resource_string(__name__, path).decode("utf-8")
                for path in paths]
        return os.linesep.join(strs)

    def _save_file(self, dest, content):
        if not os.path.isabs(dest):
            dest = os.path.join(self.api.app.root_path, dest)
        with codecs.open(dest, "w", encoding="utf-8") as f:
            f.write(content)

    def resjs(self, dest='static/res.js'):
        """generate res.js

        :param dest: dest path
        """
        tmpl = self._read_file(
            'tmpl/res-ajax.js', 'tmpl/res-promise.js', 'tmpl/res-core.js')
        apiinfo = json.dumps(self.data, ensure_ascii=False, indent=4)
        rendered = Template(tmpl).render(apiinfo=apiinfo)
        self._save_file(dest, rendered)

    def resdocs(self, dest='static/resdocs.html', resjs='/static/res.js',
                bootstrap='http://cdn.bootcss.com/bootstrap/3.3.6/css/bootstrap.min.css'):
        """generate resdocs.html

        :param dest: dest path
        :param resjs: res.js file's path
        :param bootstrap: bootstrap.css file's path
        """
        tmpl = self._read_file('tmpl/resdocs.html')
        rendered = Template(tmpl).render(
            apiinfo=self.data, resjs=resjs, bootstrap=bootstrap)
        self._save_file(dest, rendered)

    def permission(self, dest='static/permission.html', resjs='/static/res.js',
                   bootstrap='http://cdn.bootcss.com/bootstrap/3.3.6/css/bootstrap.min.css',
                   vuejs='http://cdn.bootcss.com/vue/1.0.16/vue.min.js'):
        """generate permission.html

        :param dest: dest path
        :param resjs: res.js file's path
        :param bootstrap: bootstrap.css file's path
        :param vuejs: vuejs.js file's path
        """
        tmpl = self._read_file('tmpl/permission.html')
        rendered = Template(tmpl).render(
            apiinfo=self.data, resjs=resjs, bootstrap=bootstrap, vuejs=vuejs)
        self._save_file(dest, rendered)
<commit_msg>fix AttributeError: 'Action' object has no attribute '__dict__'<commit_after>#!/usr/bin/env python
# coding: utf-8
from __future__ import unicode_literals, absolute_import, print_function
from flask_restaction import Resource
from jinja2 import Template
import codecs
from pkg_resources import resource_string
import os
from flask import json

def parse_api(api):
    """parse api"""
    data = {"url_prefix": "", "auth_header": ""}
    data["docs"] = api.docs
    if api.url_prefix:
        data["url_prefix"] = api.url_prefix
    if api.auth:
        data["auth_header"] = api.auth.auth_header
    resources = []
    for name, res in api.resources.items():
        reaource = {"name": name, "docs": res["docs"]}
        actions = [dict(x._asdict()) for x in res["actions"]]
        reaource["actions"] = actions
        resources.append(reaource)
    data["resources"] = resources
    return data

class ApiInfo(Resource):
    """ApiInfo

    :param api: Api
    """
    action = {
        "action": "str",  # post_login
        "method": "str",  # post
        "url": "str",  # /api/user/login
        "endpoint": "str",  # user@login
        "docs": "str",  # docs of post_login
        "inputs": "str",  # schema_inputs[action]
        "outputs": "str",  # schema_outputs[action]
    }
    schema_outputs = {
        "get": {
            "url_prefix": "str",  # /api
            "auth_header": "str",  # Authorization
            "docs": "str",  # docs of api
            "resources": [{
                "name": "str",  # user
                "docs": "str",  # docs of user
                "actions": [action]
            }],
        }
    }

    def __init__(self, api):
        self.data = parse_api(api)

    def get(self):
        return self.data

class Gen(object):
    """Generate tools for docs res.js

    if dest param of generate methods is not absolute path,
    then join app.root_path and dest as dest path.
    :param api: Api
    """

    def __init__(self, api):
        self.api = api
        self.data = parse_api(api)
        app = self.api.app
        if not os.path.exists(app.static_folder):
            os.makedirs(app.static_folder)

    def _read_file(self, *paths):
        strs = [resource_string(__name__, path).decode("utf-8")
                for path in paths]
        return os.linesep.join(strs)

    def _save_file(self, dest, content):
        if not os.path.isabs(dest):
            dest = os.path.join(self.api.app.root_path, dest)
        with codecs.open(dest, "w", encoding="utf-8") as f:
            f.write(content)

    def resjs(self, dest='static/res.js'):
        """generate res.js

        :param dest: dest path
        """
        tmpl = self._read_file(
            'tmpl/res-ajax.js', 'tmpl/res-promise.js', 'tmpl/res-core.js')
        apiinfo = json.dumps(self.data, ensure_ascii=False, indent=4)
        rendered = Template(tmpl).render(apiinfo=apiinfo)
        self._save_file(dest, rendered)

    def resdocs(self, dest='static/resdocs.html', resjs='/static/res.js',
                bootstrap='http://cdn.bootcss.com/bootstrap/3.3.6/css/bootstrap.min.css'):
        """generate resdocs.html

        :param dest: dest path
        :param resjs: res.js file's path
        :param bootstrap: bootstrap.css file's path
        """
        tmpl = self._read_file('tmpl/resdocs.html')
        rendered = Template(tmpl).render(
            apiinfo=self.data, resjs=resjs, bootstrap=bootstrap)
        self._save_file(dest, rendered)

    def permission(self, dest='static/permission.html', resjs='/static/res.js',
                   bootstrap='http://cdn.bootcss.com/bootstrap/3.3.6/css/bootstrap.min.css',
                   vuejs='http://cdn.bootcss.com/vue/1.0.16/vue.min.js'):
        """generate permission.html

        :param dest: dest path
        :param resjs: res.js file's path
        :param bootstrap: bootstrap.css file's path
        :param vuejs: vuejs.js file's path
        """
        tmpl = self._read_file('tmpl/permission.html')
        rendered = Template(tmpl).render(
            apiinfo=self.data, resjs=resjs, bootstrap=bootstrap, vuejs=vuejs)
        self._save_file(dest, rendered)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf8 -*-
################################################################################
##
## Copyright (C) 2012 Typhos
##
## This Source Code Form is subject to the terms of the Mozilla Public
## License, v. 2.0. If a copy of the MPL was not distributed with this
## file, You can obtain one at http://mozilla.org/MPL/2.0/.
##
################################################################################

import argparse
import time

import bplib

### Emote loading

class EmoteFile:
    def __init__(self, meta, emotes):
        self.meta = meta
        self.emotes = emotes

        self.name = meta["Name"]

    @classmethod
    def from_data(cls, data):
        meta = data.pop("Meta")
        emote_data = data.pop("Emotes")
        emotes = {}

        def update(d):
            for name_pair in d:
                if name_pair in emotes:
                    print("ERROR: Duplicate emote", name_pair)
            emotes.update(d)

        update(convert_spritesheets(emote_data.pop("Spritesheets", {})))
        update(convert_customs(emote_data.pop("Custom", {})))

        for key in emote_data:
            print("ERROR: Unknown emote section in file")
        for key in data:
            print("ERROR: Unknown key in file")

        return cls(meta, emotes)

    def __repr__(self):
        return "EmoteFile(%r, %r)" % (self.meta, self.emotes)

def convert_spritesheets(raw_spritesheets):
    # Loads all spritesheets
    emotes = {}
    for (image_url, emote_list) in raw_spritesheets.items():
        for data in emote_list:
            emote = bplib.Emote.from_data(image_url, data)
            name_pair = (emote.name, emote.suffix)
            if name_pair in emotes:
                print("ERROR: Duplicate emotes", name_pair)
            emotes[name_pair] = emote
    return emotes

def convert_customs(customs):
    # Loads all custom emotes
    emotes = {}
    for data in customs:
        emote = bplib.RawEmote.from_data(data)
        name_pair = (emote.name, emote.suffix)
        if name_pair in emotes:
            print("ERROR: Duplicate emotes", name_pair)
        emotes[name_pair] = emote
    return emotes

### Directive application

def resolve_emotes(files, data):
    file_ignores = data.pop("File Ignores")
    sorting_rules = data.pop("Sorting")
    _conflict_rules = data.pop("Conflicts")
    directive_map = data.pop("Directives")

    conflict_rules = {tuple(d["Name"]): d["Favor"] for d in _conflict_rules}

    def overrides(base, new, name_pair):
        for rule in sorting_rules:
            if rule == [base, "*"]: # Old subreddit wins all
                return False
            elif rule == [new, "*"]: # New subreddit wins all
                return True
            elif rule == [base, new]: # Old subreddit wins vs. new
                return False
            elif rule == [new, base]: # New subreddit wins vs. old
                return True

        if name_pair in conflict_rules:
            # Does either subreddit explicitly win for this emote?
            favor = conflict_rules[name_pair]
            if favor == base:
                return False
            elif favor == new:
                return True

        # No solution
        return None

    for filename in file_ignores:
        if filename in files:
            del files[filename]

    # Converts a list of files into one big emote map
    emotes = {}
    sources = {}

    conflicts = {}

    # Sort all emotes by prioritization
    for file in files.values():
        for (name_pair, emote) in file.emotes.items():
            if name_pair in emotes:
                result = overrides(sources[name_pair], file.name, name_pair)
                if result:
                    # Get rid of the old one so it can't hurt anyone anymore
                    del files[sources[name_pair]].emotes[name_pair]
                    if name_pair in conflicts:
                        del conflicts[name_pair]
                    # Replace emote
                    emotes[name_pair] = emote
                    sources[name_pair] = file.name
                elif result is None:
                    # ?!? first file wins I guess
                    conflicts[name_pair] = (file.name, sources[name_pair])
            else:
                emotes[name_pair] = emote
                sources[name_pair] = file.name

    for (name_pair, (old, new)) in conflicts.items():
        print("ERROR: CONFLICT between %s and %s over %s" % (old, new, name_pair))

    # Apply custom directives
    for (filename, directives) in directive_map.items():
        for d in directives:
            if isinstance(d, str):
                op = d
                args = ()
            else:
                op = d[0]
                args = d[1:]
            Directives[op](emotes, files[filename], *args)

    return emotes

def d_add_css(emotes, file, emote, css):
    emotes[parse_emote(emote)].css.update(css)

def d_mark_nsfw(emotes, file, *es):
    for e in es:
        file.emotes[parse_emote(e)].nsfw = True

def d_set_selector(emotes, file, emote, to):
    assert file.emotes[parse_emote(emote)].selector is None
    file.emotes[parse_emote(emote)].selector = to

def d_disable_css(emotes, file, emote):
    file.emotes[parse_emote(emote)].nocss = True

def d_remove_emote(emotes, file, emote):
    del emotes[parse_emote(emote)]
    del file.emotes[parse_emote(emote)]

def parse_emote(s):
    if isinstance(s, list):
        return tuple(s)
    else:
        return (s, None)

Directives = {
    "AddCSS": d_add_css,
    "MarkNSFW": d_mark_nsfw,
    "SetSelector": d_set_selector,
    "DisableCSS": d_disable_css,
    "Remove": d_remove_emote,
    }

### Generation

def build_css(emotes):
    css_rules, nsfw_css_rules = {}, {}

    for emote in emotes.values():
        if emote.nocss:
            continue

        target = (nsfw_css_rules if emote.nsfw else css_rules)
        selector, properties = emote.to_css()
        if selector in target:
            print("ERROR: Selector %r used twice!" % (selector))
        target[selector] = properties

    return css_rules, nsfw_css_rules

def build_js(emotes):
    emote_map = {}

    for emote in emotes.values():
        if emote.suffix is not None:
            # Assume we should ignore- I can't think of when this isn't true
            continue

        assert emote.name not in emote_map
        emote_map[emote.name] = int(emote.nsfw) + 1

    return emote_map

def condense_css(rules):
    # Locate all known CSS properties, and sort selectors by their value
    properties = {}
    for (selector, props) in rules.items():
        for (prop_name, prop_value) in props.items():
            properties.setdefault(prop_name, {}).setdefault(prop_value, [])
            properties[prop_name][prop_value].append(selector)

    def common_selectors(props):
        i = iter(props.items())

        try:
            (k, v) = next(i)
        except StopIteration:
            return set()
        common = set(properties.get(k, {}).get(v, []))

        for (k, v) in i:
            these = set(properties.get(k, {}).get(v, []))
            common = common.intersection(these)

        return common

    def condense(common_props):
        # Assume they're common
        selectors = common_selectors(common_props)
        if len(selectors) <= 1:
            return

        # TODO: frozenset? probably doesn't matter if we sort
        sel_string = ",".join(sorted(selectors))

        props_chars = sum((len(key) + len(val) + 1) for (key, val) in common_props.items())
        chars_added = len(sel_string) + props_chars

        chars_removed = 0
        for sel in selectors:
            chars_removed += props_chars
            if len(rules[sel]) <= len(common_props):
                chars_removed += len(sel) + 2

        if chars_added > chars_removed:
            return

        existing_props = rules.setdefault(sel_string, {})
        for (prop_name, value) in common_props.items():
            # Make sure we're not screwing anything up
            if prop_name in existing_props:
                print("WARNING: condensing same property twice", prop_name, "=", value)
                assert existing_props[prop_name] == value

            # Write property to rule
            existing_props[prop_name] = value
            properties[prop_name][value].append(sel_string)
            for selector in selectors:
                properties[prop_name][value].remove(selector)

            # Delete property from old rules
            for selector in selectors:
                rules[selector].pop(prop_name)
                # If it's empty, remove it entirely
                if not rules[selector]:
                    del rules[selector]

    # TODO: Would be nice to automatically seek out stuff we can efficiently
    # collapse, but for now, this achieves great gains for little complexity.

    # Remove all useless background-position's
    for selector in properties.get("background-position", {}).get("0px 0px", []):
        rules[selector].pop("background-position")
        if not rules[selector]:
            del rules[selector]
    if "0px 0px" in properties["background-position"]:
        del properties["background-position"]["0px 0px"]

    # Condense multi-emote spritesheets (probably most of our savings)
    for (image_url, selectors) in list(properties.get("background-image", {}).items()):
        if len(selectors) > 1:
            # For some reason, condensing these properties here gains more savings
            # than doing them separately. Oh well.
            condense({"background-image": image_url, "display": "block", "clear": "none", "float": "left"})
            #condense({"background-image": image_url})

    # Condense similar background-position's. Not likely to make a big difference
    # except for a few very similar spritesheet grids.
    for position in list(properties["background-position"]):
        condense({"background-position": position})

    # Condense by width/height, since many emotes have the same dimensions
    for (width, w_selectors) in [(w, s) for (w, s) in properties["width"].items() if len(s) > 1]:
        for (height, h_selectors) in [(h, s) for (h, s) in properties["height"].items() if len(s) > 1]:
            if any(s in h_selectors for s in w_selectors):
                condense({"height": height, "width": width})

    for width in properties["width"]:
        condense({"width": width})

    for height in properties["height"]:
        condense({"height": height})

    # Locate and combine identical rules
    for (selector, props) in rules.copy().items():
        condense(props.copy())

AutogenHeader = """
/*
 * This file is AUTOMATICALLY GENERATED. DO NOT EDIT.
 * Generated at %s.
 */

""" % (time.strftime("%c"))

def dump_css(file, rules):
    file.write(AutogenHeader)
    for (selector, properties) in rules.items():
        property_strings = ["%s:%s" % i for i in properties.items()]
        s = "%s{%s}\n" % (selector, ";".join(property_strings))
        file.write(s)

def dump_js(file, js_map):
    file.write(AutogenHeader)
    file.write("var emote_map = {\n")

    strings = ["%r:%r" % i for i in js_map.items()]
    file.write(",\n".join(strings))

    file.write("\n}\n")

### Main

def main():
    parser = argparse.ArgumentParser(description="Generate addon data files from emotes")
    parser.add_argument("-j", "--js", help="Output JS file", default="build/emote-map.js")
    parser.add_argument("-c", "--css", help="Output CSS file", default="build/emote-classes.css")
    parser.add_argument("-n", "--nsfw-css", help="Output NSFW CSS file", default="build/nsfw-emote-classes.css")
    parser.add_argument("-d", "--directives", help="Processing directives",
                        default="data/emote-directives.yaml", type=argparse.FileType("r"))
    parser.add_argument("emotes", help="Input emote files", nargs="+")
    args = parser.parse_args()

    files = {}

    print("Loading emotes")
    for (i, filename) in enumerate(args.emotes):
        with open(filename) as file:
            data = bplib.load_yaml_file(file)
            efile = EmoteFile.from_data(data)
            files[efile.meta["Name"]] = efile

    print("Processing")
    emotes = resolve_emotes(files, bplib.load_yaml_file(args.directives))
    css_rules, nsfw_css_rules = build_css(emotes)
    js_map = build_js(emotes)
    condense_css(css_rules)
    condense_css(nsfw_css_rules)

    print("Dumping")
    with open(args.css, "w") as file:
        dump_css(file, css_rules)
    with open(args.nsfw_css, "w") as file:
        dump_css(file, nsfw_css_rules)
    with open(args.js, "w") as file:
        dump_js(file, js_map)

if __name__ == "__main__":
    main()
<commit_msg>More assert's to ensure the CSS remains valid<commit_after>#!/usr/bin/env python3
# -*- coding: utf8 -*-
################################################################################
##
## Copyright (C) 2012 Typhos
##
## This Source Code Form is subject to the terms of the Mozilla Public
## License, v. 2.0. If a copy of the MPL was not distributed with this
## file, You can obtain one at http://mozilla.org/MPL/2.0/.
##
################################################################################

import argparse
import time

import bplib

### Emote loading

class EmoteFile:
    def __init__(self, meta, emotes):
        self.meta = meta
        self.emotes = emotes

        self.name = meta["Name"]

    @classmethod
    def from_data(cls, data):
        meta = data.pop("Meta")
        emote_data = data.pop("Emotes")
        emotes = {}

        def update(d):
            for name_pair in d:
                if name_pair in emotes:
                    print("ERROR: Duplicate emote", name_pair)
            emotes.update(d)

        update(convert_spritesheets(emote_data.pop("Spritesheets", {})))
        update(convert_customs(emote_data.pop("Custom", {})))

        for key in emote_data:
            print("ERROR: Unknown emote section in file")
        for key in data:
            print("ERROR: Unknown key in file")

        return cls(meta, emotes)

    def __repr__(self):
        return "EmoteFile(%r, %r)" % (self.meta, self.emotes)

def convert_spritesheets(raw_spritesheets):
    # Loads all spritesheets
    emotes = {}
    for (image_url, emote_list) in raw_spritesheets.items():
        for data in emote_list:
            emote = bplib.Emote.from_data(image_url, data)
            name_pair = (emote.name, emote.suffix)
            if name_pair in emotes:
                print("ERROR: Duplicate emotes", name_pair)
            emotes[name_pair] = emote
    return emotes

def convert_customs(customs):
    # Loads all custom emotes
    emotes = {}
    for data in customs:
        emote = bplib.RawEmote.from_data(data)
        name_pair = (emote.name, emote.suffix)
        if name_pair in emotes:
            print("ERROR: Duplicate emotes", name_pair)
        emotes[name_pair] = emote
    return emotes

### Directive application

def resolve_emotes(files, data):
    file_ignores = data.pop("File Ignores")
    sorting_rules = data.pop("Sorting")
    _conflict_rules = data.pop("Conflicts")
    directive_map = data.pop("Directives")

    conflict_rules = {tuple(d["Name"]): d["Favor"] for d in _conflict_rules}

    def overrides(base, new, name_pair):
        for rule in sorting_rules:
            if rule == [base, "*"]: # Old subreddit wins all
                return False
            elif rule == [new, "*"]: # New subreddit wins all
                return True
            elif rule == [base, new]: # Old subreddit wins vs. new
                return False
            elif rule == [new, base]: # New subreddit wins vs. old
                return True

        if name_pair in conflict_rules:
            # Does either subreddit explicitly win for this emote?
            favor = conflict_rules[name_pair]
            if favor == base:
                return False
            elif favor == new:
                return True

        # No solution
        return None

    for filename in file_ignores:
        if filename in files:
            del files[filename]

    # Converts a list of files into one big emote map
    emotes = {}
    sources = {}

    conflicts = {}

    # Sort all emotes by prioritization
    for file in files.values():
        for (name_pair, emote) in file.emotes.items():
            if name_pair in emotes:
                result = overrides(sources[name_pair], file.name, name_pair)
                if result:
                    # Get rid of the old one so it can't hurt anyone anymore
                    del files[sources[name_pair]].emotes[name_pair]
                    if name_pair in conflicts:
                        del conflicts[name_pair]
                    # Replace emote
                    emotes[name_pair] = emote
                    sources[name_pair] = file.name
                elif result is None:
                    # ?!? first file wins I guess
                    conflicts[name_pair] = (file.name, sources[name_pair])
            else:
                emotes[name_pair] = emote
                sources[name_pair] = file.name

    for (name_pair, (old, new)) in conflicts.items():
        print("ERROR: CONFLICT between %s and %s over %s" % (old, new, name_pair))

    # Apply custom directives
    for (filename, directives) in directive_map.items():
        for d in directives:
            if isinstance(d, str):
                op = d
                args = ()
            else:
                op = d[0]
                args = d[1:]
            Directives[op](emotes, files[filename], *args)

    return emotes

def d_add_css(emotes, file, emote, css):
    emotes[parse_emote(emote)].css.update(css)

def d_mark_nsfw(emotes, file, *es):
    for e in es:
        file.emotes[parse_emote(e)].nsfw = True

def d_set_selector(emotes, file, emote, to):
    assert file.emotes[parse_emote(emote)].selector is None
    file.emotes[parse_emote(emote)].selector = to

def d_disable_css(emotes, file, emote):
    file.emotes[parse_emote(emote)].nocss = True

def d_remove_emote(emotes, file, emote):
    del emotes[parse_emote(emote)]
    del file.emotes[parse_emote(emote)]

def parse_emote(s):
    if isinstance(s, list):
        return tuple(s)
    else:
        return (s, None)

Directives = {
    "AddCSS": d_add_css,
    "MarkNSFW": d_mark_nsfw,
    "SetSelector": d_set_selector,
    "DisableCSS": d_disable_css,
    "Remove": d_remove_emote,
    }

### Generation

def build_css(emotes):
    css_rules, nsfw_css_rules = {}, {}

    for emote in emotes.values():
        if emote.nocss:
            continue

        target = (nsfw_css_rules if emote.nsfw else css_rules)
        selector, properties = emote.to_css()
        if selector in target:
            print("ERROR: Selector %r used twice!" % (selector))
        target[selector] = properties

    return css_rules, nsfw_css_rules

def build_js(emotes):
    emote_map = {}

    for emote in emotes.values():
        if emote.suffix is not None:
            # Assume we should ignore- I can't think of when this isn't true
            continue

        assert emote.name not in emote_map
        emote_map[emote.name] = int(emote.nsfw) + 1

    return emote_map

def condense_css(rules):
    # Make a copy for validation purposes
    base_rules = {sel: props.copy() for (sel, props) in rules.items()}

    # Locate all known CSS properties, and sort selectors by their value
    properties = {}
    for (selector, props) in rules.items():
        for (prop_name, prop_value) in props.items():
            properties.setdefault(prop_name, {}).setdefault(prop_value, [])
            properties[prop_name][prop_value].append(selector)

    def common_selectors(props):
        i = iter(props.items())

        try:
            (k, v) = next(i)
        except StopIteration:
            return set()
        common = set(properties.get(k, {}).get(v, []))

        for (k, v) in i:
            these = set(properties.get(k, {}).get(v, []))
            common = common.intersection(these)

        return common

    def condense(common_props):
        # Assume they're common
        selectors = common_selectors(common_props)
        if len(selectors) <= 1:
            return

        # len("property:value") for each property + semicolons
        props_chars = sum((len(key) + len(val) + 1) for (key, val) in common_props.items())

        # TODO: frozenset? probably doesn't matter if we sort
        sel_string = ",".join(sorted(selectors))

        chars_added = len(sel_string) + props_chars

        chars_removed = 0
        for sel in selectors:
            chars_removed += props_chars
            if len(rules[sel]) <= len(common_props):
                chars_removed += len(sel) + 2

        if chars_added > chars_removed:
            return

        existing_props = rules.setdefault(sel_string, {})
        if existing_props:
            # Ensure compatibility
            for (prop_name, value) in existing_props.items():
                # Don't overwrite anything not permitted
                if prop_name in common_props:
                    assert common_props[prop_name] == value

                # Don't bring anything in not expected
                for sel in selectors:
                    # Also fails if the property didn't exist before
                    assert base_rules[sel][prop_name] == value

        for (prop_name, value) in common_props.items():
            # Make sure we're not screwing anything up
            if prop_name in existing_props:
                print("WARNING: condensing same property twice", prop_name, "=", value)
                assert existing_props[prop_name] == value

            # Write property to rule
            existing_props[prop_name] = value
            properties[prop_name][value].append(sel_string)
            for selector in selectors:
                properties[prop_name][value].remove(selector)

            # Delete property from old rules
            for selector in selectors:
                rules[selector].pop(prop_name)
                # If it's empty, remove it entirely
                if not rules[selector]:
                    del rules[selector]

    # TODO: Would be nice to automatically seek out stuff we can efficiently
    # collapse, but for now, this achieves great gains for little complexity.

    # Remove all useless background-position's
    for selector in properties.get("background-position", {}).get("0px 0px", []):
        del rules[selector]["background-position"]
        if not rules[selector]:
            del rules[selector]
    if "0px 0px" in properties["background-position"]:
        del properties["background-position"]["0px 0px"]

    # Condense multi-emote spritesheets (probably most of our savings)
    for (image_url, selectors) in list(properties.get("background-image", {}).items()):
        if len(selectors) > 1:
            # For some reason, condensing these properties here gains more savings
            # than doing them separately. Oh well.
            condense({"background-image": image_url, "display": "block", "clear": "none", "float": "left"})

    # Condense similar background-position's. Not likely to make a big difference
    # except for a few very similar spritesheet grids.
    for position in list(properties["background-position"]):
        condense({"background-position": position})

    # Condense by width/height, since many emotes have the same dimensions
    for (width, w_selectors) in [(w, s) for (w, s) in properties["width"].items() if len(s) > 1]:
        for (height, h_selectors) in [(h, s) for (h, s) in properties["height"].items() if len(s) > 1]:
            if set(h_selectors).intersection(w_selectors): # Any in common? Try condensing the pair
                condense({"height": height, "width": width})

    for width in properties["width"]:
        condense({"width": width})

    for height in properties["height"]:
        condense({"height": height})

    # Locate and combine identical rules
    for (selector, props) in rules.copy().items():
        condense(props.copy())

AutogenHeader = """
/*
 * This file is AUTOMATICALLY GENERATED. DO NOT EDIT.
 * Generated at %s.
 */

""" % (time.strftime("%c"))

def dump_css(file, rules):
    file.write(AutogenHeader)
    for (selector, properties) in rules.items():
        property_strings = ["%s:%s" % i for i in properties.items()]
        s = "%s{%s}\n" % (selector, ";".join(property_strings))
        file.write(s)

def dump_js(file, js_map):
    file.write(AutogenHeader)
    file.write("var emote_map = {\n")

    strings = ["%r:%r" % i for i in js_map.items()]
    file.write(",\n".join(strings))

    file.write("\n}\n")

### Main

def main():
    parser = argparse.ArgumentParser(description="Generate addon data files from emotes")
    parser.add_argument("-j", "--js", help="Output JS file", default="build/emote-map.js")
    parser.add_argument("-c", "--css", help="Output CSS file", default="build/emote-classes.css")
    parser.add_argument("-n", "--nsfw-css", help="Output NSFW CSS file", default="build/nsfw-emote-classes.css")
    parser.add_argument("-d", "--directives", help="Processing directives",
                        default="data/emote-directives.yaml", type=argparse.FileType("r"))
    parser.add_argument("emotes", help="Input emote files", nargs="+")
    args = parser.parse_args()

    files = {}

    print("Loading emotes")
    for (i, filename) in enumerate(args.emotes):
        with open(filename) as file:
            data = bplib.load_yaml_file(file)
            efile = EmoteFile.from_data(data)
            files[efile.meta["Name"]] = efile

    print("Processing")
    emotes = resolve_emotes(files, bplib.load_yaml_file(args.directives))
    css_rules, nsfw_css_rules = build_css(emotes)
    js_map = build_js(emotes)
    condense_css(css_rules)
    condense_css(nsfw_css_rules)

    print("Dumping")
    with open(args.css, "w") as file:
        dump_css(file, css_rules)
    with open(args.nsfw_css, "w") as file:
        dump_css(file, nsfw_css_rules)
    with open(args.js, "w") as file:
        dump_js(file, js_map)

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


<commit_before>import subprocess

import os

class PxFile(object):
    def __init__(self):
        self._device_number = None

    def __repr__(self):
        # I guess this is really what __str__ should be doing, but the point of
        # implementing this method is to make the py.test output more readable,
        # and py.test calls repr() and not str().
        return str(self.pid) + ":" + self.name

def device_to_number(device):
    if device is None:
        return None

    number = int(device, 16)
    if number == 0:
        # 0x000lotsofmore000 is what we get on lsof 4.86 and Linux 4.2.0
        # when lsof doesn't have root privileges
        return None

    return number

def call_lsof():
    """
    Call lsof and return the result as one big string
    """
    env = os.environ.copy()
    if "LANG" in env:
        del env["LANG"]

    # See OUTPUT FOR OTHER PROGRAMS: http://linux.die.net/man/8/lsof
    # Output lines can be in one of two formats:
    # 1. "pPID@" (with @ meaning NUL)
    # 2. "fFD@aACCESSMODE@tTYPE@nNAME@"
    lsof = subprocess.Popen(["lsof", '-F', 'fnaptd0'],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            env=env)
    return lsof.communicate()[0]

def lsof_to_files(lsof):
    pid = None
    file = None
    files = []
    for shard in lsof.split('\0'):
        if shard[0] == "\n":
            # Some shards start with newlines. Looks pretty when viewing the
            # lsof output in less, but makes the parsing code have to deal with
            # it.
            shard = shard[1:]

        if len(shard) == 0:
            # The output ends with a single newline, which we just stripped away
            break

        type = shard[0]
        value = shard[1:]

        if type == 'p':
            pid = int(value)
        elif type == 'f':
            file = PxFile()
            file.fd = value
            file.pid = pid
            file.type = "??"
            file.device = None
            file.device_number = None
            files.append(file)
        elif type == 'a':
            file.access = {
                ' ': None,
                'r': "r",
                'w': "w",
                'u': "rw"}[value]
        elif type == 't':
            file.type = value
        elif type == 'd':
            file.device = value
            file.device_number = device_to_number(value)
        elif type == 'n':
            file.name = value
            file.plain_name = value
            if file.type != "REG":
                # Decorate non-regular files with their type
                file.name = "[" + file.type + "] " + file.name

        else:
            raise Exception("Unhandled type <{}> for shard <{}>".format(type, shard))

    return files

def get_all():
    return lsof_to_files(call_lsof())
<commit_msg>Make lsof_to_files() about 15% faster<commit_after>import subprocess

import os

class PxFile(object):
    def __init__(self):
        self._device_number = None

    def __repr__(self):
        # I guess this is really what __str__ should be doing, but the point of
        # implementing this method is to make the py.test output more readable,
        # and py.test calls repr() and not str().
        return str(self.pid) + ":" + self.name

def device_to_number(device):
    if device is None:
        return None

    number = int(device, 16)
    if number == 0:
        # 0x000lotsofmore000 is what we get on lsof 4.86 and Linux 4.2.0
        # when lsof doesn't have root privileges
        return None

    return number

def call_lsof():
    """
    Call lsof and return the result as one big string
    """
    env = os.environ.copy()
    if "LANG" in env:
        del env["LANG"]

    # See OUTPUT FOR OTHER PROGRAMS: http://linux.die.net/man/8/lsof
    # Output lines can be in one of two formats:
    # 1. "pPID@" (with @ meaning NUL)
    # 2. "fFD@aACCESSMODE@tTYPE@nNAME@"
    lsof = subprocess.Popen(["lsof", '-F', 'fnaptd0'],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            env=env)
    return lsof.communicate()[0]

def lsof_to_files(lsof):
    pid = None
    file = None
    files = []
    for shard in lsof.split('\0'):
        if shard[0] == "\n":
            # Some shards start with newlines. Looks pretty when viewing the
            # lsof output in less, but makes the parsing code have to deal with
            # it.
            shard = shard[1:]

        if not shard:
            # The output ends with a single newline, which we just stripped away
            break

        type = shard[0]
        value = shard[1:]

        if type == 'p':
            pid = int(value)
        elif type == 'f':
            file = PxFile()
            file.fd = value
            file.pid = pid
            file.type = "??"
            file.device = None
            file.device_number = None
            files.append(file)
        elif type == 'a':
            file.access = {
                ' ': None,
                'r': "r",
                'w': "w",
                'u': "rw"}[value]
        elif type == 't':
            file.type = value
        elif type == 'd':
            file.device = value
            file.device_number = device_to_number(value)
        elif type == 'n':
            file.name = value
            file.plain_name = value
            if file.type != "REG":
                # Decorate non-regular files with their type
                file.name = "[" + file.type + "] " + file.name

        else:
            raise Exception("Unhandled type <{}> for shard <{}>".format(type, shard))

    return files

def get_all():
    return lsof_to_files(call_lsof())
<|endoftext|>


<commit_before>import asyncio
import calendar
import datetime
import struct
from bson import objectid

from vj4 import app
from vj4 import constant
from vj4 import error
from vj4.handler import base
from vj4.model import builtin
from vj4.model import document
from vj4.model import domain
from vj4.model import fs
from vj4.model import record
from vj4.model import user
from vj4.model.adaptor import contest
from vj4.model.adaptor import problem
from vj4.service import bus
from vj4.util import options

@app.route('/records', 'record_main')
class RecordMainHandler(base.Handler):
  @base.get_argument
  @base.sanitize
  async def get(self, *, uid_or_name: str='', pid: str='', tid: str=''):
    query = {}
    if uid_or_name:
      try:
        query['uid'] = int(uid_or_name)
      except ValueError:
        udoc = await user.get_by_uname(uid_or_name)
        if not udoc:
          raise error.UserNotFoundError(uid_or_name) from None
        query['uid'] = udoc['_id']
    if pid:
      pid = document.convert_doc_id(pid)
      query['domain_id'] = self.domain_id
      query['pid'] = pid
    if tid:
      tid = document.convert_doc_id(tid)
      query['domain_id'] = self.domain_id
      query['tid'] = tid
    # TODO(iceboy): projection, pagination.
    rdocs = await record.get_all_multi(**query,
      get_hidden=self.has_priv(builtin.PRIV_VIEW_HIDDEN_RECORD)).sort([('_id', -1)]).limit(50).to_list(None)
    # TODO(iceboy): projection.
    udict, pdict = await asyncio.gather(
        user.get_dict(rdoc['uid'] for rdoc in rdocs),
        problem.get_dict_multi_domain((rdoc['domain_id'], rdoc['pid']) for rdoc in rdocs))
    # statistics
    statistics = None
    if self.has_priv(builtin.PRIV_VIEW_JUDGE_STATISTICS):
      ts = calendar.timegm(datetime.datetime.utcnow().utctimetuple())
      day_count, week_count, month_count, year_count, rcount = await asyncio.gather(
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 7 * 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 30 * 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - int(365.2425 * 24 * 3600)) + struct.pack('b', -1) * 8)),
          record.get_count())
      statistics = {'day': day_count, 'week': week_count, 'month': month_count,
                    'year': year_count, 'total': rcount}
    self.render('record_main.html', rdocs=rdocs, udict=udict, pdict=pdict, statistics=statistics,
                filter_uid_or_name=uid_or_name, filter_pid=pid, filter_tid=tid)

@app.connection_route('/records-conn', 'record_main-conn')
class RecordMainConnection(base.Connection):
  async def on_open(self):
    await super(RecordMainConnection, self).on_open()
    bus.subscribe(self.on_record_change, ['record_change'])

  async def on_record_change(self, e):
    rdoc = await record.get(e['value'], record.PROJECTION_PUBLIC)
    # check permission for visibility: contest
    if rdoc['tid']:
      now = datetime.datetime.utcnow()
      tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
      if (not contest.RULES[tdoc['rule']].show_func(tdoc, now)
          and (self.domain_id != tdoc['domain_id']
               or not self.has_perm(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS))):
        return
    # TODO(iceboy): join from event to improve performance?
    # TODO(iceboy): projection.
    udoc, pdoc = await asyncio.gather(user.get_by_uid(rdoc['uid']),
                                      problem.get(rdoc['domain_id'], rdoc['pid']))
    # check permission for visibility: hidden problem
    if pdoc.get('hidden', False) and (pdoc['domain_id'] != self.domain_id
                                      or not self.has_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)):
      pdoc = None
    self.send(html=self.render_html('record_main_tr.html', rdoc=rdoc, udoc=udoc, pdoc=pdoc))

  async def on_close(self):
    bus.unsubscribe(self.on_record_change)

@app.route('/records/{rid}', 'record_detail')
class RecordDetailHandler(base.Handler):
  @base.route_argument
  @base.sanitize
  async def get(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if not rdoc:
      raise error.RecordNotFoundError(rid)
    # TODO(iceboy): Check domain permission, permission for visibility in place.
    if rdoc['domain_id'] != self.domain_id:
      self.redirect(self.reverse_url('record_detail', rid=rid, domain_id=rdoc['domain_id']))
      return
    # check permission for visibility: contest
    show_status = True
    if rdoc['tid']:
      now = datetime.datetime.utcnow()
      try:
        tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
        show_status = contest.RULES[tdoc['rule']].show_func(tdoc, now) \
                      or self.has_perm(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS)
      except error.DocumentNotFoundError:
        tdoc = None
    else:
      tdoc = None
    # TODO(twd2): futher check permission for visibility.
    if (not self.own(rdoc, field='uid')
        and not self.has_perm(builtin.PERM_READ_RECORD_CODE)
        and not self.has_priv(builtin.PRIV_READ_RECORD_CODE)):
      del rdoc['code']
    if not show_status and 'code' not in rdoc:
      raise error.PermissionError(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS)
    udoc, dudoc = await asyncio.gather(
        user.get_by_uid(rdoc['uid']),
        domain.get_user(self.domain_id, rdoc['uid']))
    try:
      pdoc = await problem.get(rdoc['domain_id'], rdoc['pid'])
    except error.ProblemNotFoundError:
      pdoc = {}
    if show_status and 'judge_uid' in rdoc:
      judge_udoc = await user.get_by_uid(rdoc['judge_uid'])
    else:
      judge_udoc = None
    # check permission for visibility: hidden problem
    if pdoc.get('hidden', False) and not self.has_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN):
      pdoc = None
    self.render('record_detail.html', rdoc=rdoc, udoc=udoc, dudoc=dudoc, pdoc=pdoc, tdoc=tdoc,
                judge_udoc=judge_udoc, show_status=show_status)

@app.connection_route('/records/{rid}/conn', 'record_detail-conn')
class RecordDetailConnection(base.Connection):
  async def on_open(self):
    await super(RecordDetailConnection, self).on_open()
    self.rid = objectid.ObjectId(self.request.match_info['rid'])
    bus.subscribe(self.on_record_change, ['record_change'])

  async def on_record_change(self, e):
    if e['value'] != self.rid:
      return
    rdoc = await record.get(self.rid, record.PROJECTION_PUBLIC)
    # check permission for visibility: contest
    if rdoc['tid']:
      now = datetime.datetime.utcnow()
      tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
      if (not contest.RULES[tdoc['rule']].show_func(tdoc, now)
          and (self.domain_id != tdoc['domain_id']
               or not self.has_perm(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS))):
        return
    self.send(status_html=self.render_html('record_detail_status.html', rdoc=rdoc),
              summary_html=self.render_html('record_detail_summary.html', rdoc=rdoc))

  async def on_close(self):
    bus.unsubscribe(self.on_record_change)

@app.route('/records/{rid}/rejudge', 'record_rejudge')
class RecordRejudgeHandler(base.Handler):
  @base.route_argument
  @base.post_argument
  @base.require_csrf_token
  @base.sanitize
  async def post(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if rdoc['domain_id'] == self.domain_id:
      self.check_perm(builtin.PERM_REJUDGE)
    else:
      self.check_priv(builtin.PRIV_REJUDGE)
    await record.rejudge(rdoc['_id'])
    self.json_or_redirect(self.referer_or_main)

@app.route('/records/{rid}/data', 'record_pretest_data')
class RecordPretestDataHandler(base.Handler):
  @base.route_argument
  @base.sanitize
  async def get(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if not rdoc or rdoc['type'] != constant.record.TYPE_PRETEST:
      raise error.RecordNotFoundError(rid)
    if not self.own(rdoc, builtin.PRIV_READ_PRETEST_DATA_SELF, 'uid'):
      self.check_priv(builtin.PRIV_READ_PRETEST_DATA)
    if not rdoc.get('data_id'):
      raise error.RecordDataNotFoundError(rdoc['_id'])
    secret = await fs.get_secret(rdoc['data_id'])
    if not secret:
      raise error.RecordDataNotFoundError(rdoc['_id'])
    self.redirect(options.cdn_prefix.rstrip('/') + \
                  self.reverse_url('fs_get', domain_id=builtin.DOMAIN_ID_SYSTEM,
                                   secret=secret))
<commit_msg>grumble<commit_after>import asyncio
import calendar
import datetime
import struct
from bson import objectid

from vj4 import app
from vj4 import constant
from vj4 import error
from vj4.handler import base
from vj4.model import builtin
from vj4.model import document
from vj4.model import domain
from vj4.model import fs
from vj4.model import record
from vj4.model import user
from vj4.model.adaptor import contest
from vj4.model.adaptor import problem
from vj4.service import bus
from vj4.util import options

class RecordMixin:
  def tdoc_visible(self, tdoc):
    now = datetime.datetime.utcnow()
    if (not contest.RULES[tdoc['rule']].show_func(tdoc, now)
        and (self.domain_id != tdoc['domain_id']
             or not self.has_perm(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS))):
      return False
    return True

@app.route('/records', 'record_main')
class RecordMainHandler(base.Handler):
  @base.get_argument
  @base.sanitize
  async def get(self, *, uid_or_name: str='', pid: str='', tid: str=''):
    query = {}
    if uid_or_name:
      try:
        query['uid'] = int(uid_or_name)
      except ValueError:
        udoc = await user.get_by_uname(uid_or_name)
        if not udoc:
          raise error.UserNotFoundError(uid_or_name) from None
        query['uid'] = udoc['_id']
    if pid:
      pid = document.convert_doc_id(pid)
      query['domain_id'] = self.domain_id
      query['pid'] = pid
    if tid:
      tid = document.convert_doc_id(tid)
      query['domain_id'] = self.domain_id
      query['tid'] = tid
    # TODO(iceboy): projection, pagination.
    rdocs = await record.get_all_multi(**query,
      get_hidden=self.has_priv(builtin.PRIV_VIEW_HIDDEN_RECORD)).sort([('_id', -1)]).limit(50).to_list(None)
    # TODO(iceboy): projection.
    udict, pdict = await asyncio.gather(
        user.get_dict(rdoc['uid'] for rdoc in rdocs),
        problem.get_dict_multi_domain((rdoc['domain_id'], rdoc['pid']) for rdoc in rdocs))
    # statistics
    statistics = None
    if self.has_priv(builtin.PRIV_VIEW_JUDGE_STATISTICS):
      ts = calendar.timegm(datetime.datetime.utcnow().utctimetuple())
      day_count, week_count, month_count, year_count, rcount = await asyncio.gather(
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 7 * 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - 30 * 24 * 3600) + struct.pack('b', -1) * 8)),
          record.get_count(objectid.ObjectId(
              struct.pack('>i', ts - int(365.2425 * 24 * 3600)) + struct.pack('b', -1) * 8)),
          record.get_count())
      statistics = {'day': day_count, 'week': week_count, 'month': month_count,
                    'year': year_count, 'total': rcount}
    self.render('record_main.html', rdocs=rdocs, udict=udict, pdict=pdict, statistics=statistics,
                filter_uid_or_name=uid_or_name, filter_pid=pid, filter_tid=tid)

@app.connection_route('/records-conn', 'record_main-conn')
class RecordMainConnection(base.Connection, RecordMixin):
  async def on_open(self):
    await super(RecordMainConnection, self).on_open()
    bus.subscribe(self.on_record_change, ['record_change'])

  async def on_record_change(self, e):
    rdoc = await record.get(e['value'], record.PROJECTION_PUBLIC)
    if rdoc['tid']:
      tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
      if not self.tdoc_visible(tdoc):
        return
    # TODO(iceboy): join from event to improve performance?
    # TODO(iceboy): projection.
    udoc, pdoc = await asyncio.gather(user.get_by_uid(rdoc['uid']),
                                      problem.get(rdoc['domain_id'], rdoc['pid']))
    # check permission for visibility: hidden problem
    if pdoc.get('hidden', False) and (pdoc['domain_id'] != self.domain_id
                                      or not self.has_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN)):
      pdoc = None
    self.send(html=self.render_html('record_main_tr.html', rdoc=rdoc, udoc=udoc, pdoc=pdoc))

  async def on_close(self):
    bus.unsubscribe(self.on_record_change)

@app.route('/records/{rid}', 'record_detail')
class RecordDetailHandler(base.Handler, RecordMixin):
  @base.route_argument
  @base.sanitize
  async def get(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if not rdoc:
      raise error.RecordNotFoundError(rid)
    # TODO(iceboy): Check domain permission, permission for visibility in place.
    if rdoc['domain_id'] != self.domain_id:
      self.redirect(self.reverse_url('record_detail', rid=rid, domain_id=rdoc['domain_id']))
      return
    show_status = True
    if rdoc['tid']:
      tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
      show_status = self.tdoc_visible(tdoc)
    else:
      tdoc = None
    # TODO(twd2): futher check permission for visibility.
    if (not self.own(rdoc, field='uid')
        and not self.has_perm(builtin.PERM_READ_RECORD_CODE)
        and not self.has_priv(builtin.PRIV_READ_RECORD_CODE)):
      del rdoc['code']
    if not show_status and 'code' not in rdoc:
      raise error.PermissionError(builtin.PERM_VIEW_CONTEST_HIDDEN_STATUS)
    udoc, dudoc = await asyncio.gather(
        user.get_by_uid(rdoc['uid']),
        domain.get_user(self.domain_id, rdoc['uid']))
    try:
      pdoc = await problem.get(rdoc['domain_id'], rdoc['pid'])
    except error.ProblemNotFoundError:
      pdoc = {}
    if show_status and 'judge_uid' in rdoc:
      judge_udoc = await user.get_by_uid(rdoc['judge_uid'])
    else:
      judge_udoc = None
    # check permission for visibility: hidden problem
    if pdoc.get('hidden', False) and not self.has_perm(builtin.PERM_VIEW_PROBLEM_HIDDEN):
      pdoc = None
    self.render('record_detail.html', rdoc=rdoc, udoc=udoc, dudoc=dudoc, pdoc=pdoc, tdoc=tdoc,
                judge_udoc=judge_udoc, show_status=show_status)

@app.connection_route('/records/{rid}/conn', 'record_detail-conn')
class RecordDetailConnection(base.Connection, RecordMixin):
  async def on_open(self):
    await super(RecordDetailConnection, self).on_open()
    self.rid = objectid.ObjectId(self.request.match_info['rid'])
    rdoc = await record.get(self.rid, record.PROJECTION_PUBLIC)
    if rdoc['tid']:
      self.tdoc = await contest.get(rdoc['domain_id'], rdoc['tid'])
    else:
      self.tdoc = None
    bus.subscribe(self.on_record_change, ['record_change'])
    await self.try_send_record(rdoc)

  async def on_record_change(self, e):
    if e['value'] != self.rid:
      return
    rdoc = await record.get(self.rid, record.PROJECTION_PUBLIC)
    await self.try_send_record(rdoc)

  async def try_send_record(self, rdoc):
    if self.tdoc:
      if not self.tdoc_visible(self.tdoc):
        return
    self.send(status_html=self.render_html('record_detail_status.html', rdoc=rdoc),
              summary_html=self.render_html('record_detail_summary.html', rdoc=rdoc))

  async def on_close(self):
    bus.unsubscribe(self.on_record_change)

@app.route('/records/{rid}/rejudge', 'record_rejudge')
class RecordRejudgeHandler(base.Handler):
  @base.route_argument
  @base.post_argument
  @base.require_csrf_token
  @base.sanitize
  async def post(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if rdoc['domain_id'] == self.domain_id:
      self.check_perm(builtin.PERM_REJUDGE)
    else:
      self.check_priv(builtin.PRIV_REJUDGE)
    await record.rejudge(rdoc['_id'])
    self.json_or_redirect(self.referer_or_main)

@app.route('/records/{rid}/data', 'record_pretest_data')
class RecordPretestDataHandler(base.Handler):
  @base.route_argument
  @base.sanitize
  async def get(self, *, rid: objectid.ObjectId):
    rdoc = await record.get(rid)
    if not rdoc or rdoc['type'] != constant.record.TYPE_PRETEST:
      raise error.RecordNotFoundError(rid)
    if not self.own(rdoc, builtin.PRIV_READ_PRETEST_DATA_SELF, 'uid'):
      self.check_priv(builtin.PRIV_READ_PRETEST_DATA)
    if not rdoc.get('data_id'):
      raise error.RecordDataNotFoundError(rdoc['_id'])
    secret = await fs.get_secret(rdoc['data_id'])
    if not secret:
      raise error.RecordDataNotFoundError(rdoc['_id'])
    self.redirect(options.cdn_prefix.rstrip('/') + \
                  self.reverse_url('fs_get', domain_id=builtin.DOMAIN_ID_SYSTEM,
                                   secret=secret))
<|endoftext|>


<commit_before>from flask.ext.plugins import Plugin
from flask import current_app

from flaskbb.admin.models import SettingsGroup

class FlaskBBPlugin(Plugin):

    #: Set this to true if the plugin needs to install additional things
    settings_key = None

    @property
    def installable(self):
        if self.settings_key is not None:
            return True
        return False

    @property
    def uninstallable(self):
        if self.installable:
            group = SettingsGroup.query.filter_by(key=self.settings_key).first()
            if group and len(group.settings.all()) > 0:
                return True
            return False
        return False

        # Some helpers
    def register_blueprint(self, blueprint, **kwargs):
        """Registers a blueprint."""
        current_app.register_blueprint(blueprint, **kwargs)

    def create_table(self, model, db):
        """Creates the relation for the model

        :param model: The Model which should be created
        :param db: The database instance.
        """
        if not model.__table__.exists(bind=db.engine):
            model.__table__.create(bind=db.engine)

    def drop_table(self, model, db):
        """Drops the relation for the bounded model.

        :param model: The model on which the table is bound.
        :param db: The database instance.
        """
        model.__table__.drop(bind=db.engine)

    def create_all_tables(self, models, db):
        """A interface for creating all models specified in ``models``.

        :param models: A list with models
        :param db: The database instance
        """
        for model in models:
            self.create_table(model, db)

    def drop_all_tables(self, models, db):
        """A interface for dropping all models specified in the
        variable ``models``.

        :param models: A list with models
        :param db: The database instance.
        """
        for model in models:
            self.drop_table(model, db)
<commit_msg>Renamed admin to management (1/2). Forgot to pull the latest changes from the repository..<commit_after>from flask.ext.plugins import Plugin
from flask import current_app

from flaskbb.management.models import SettingsGroup

class FlaskBBPlugin(Plugin):

    #: Set this to true if the plugin needs to install additional things
    settings_key = None

    @property
    def installable(self):
        if self.settings_key is not None:
            return True
        return False

    @property
    def uninstallable(self):
        if self.installable:
            group = SettingsGroup.query.filter_by(key=self.settings_key).first()
            if group and len(group.settings.all()) > 0:
                return True
            return False
        return False

        # Some helpers
    def register_blueprint(self, blueprint, **kwargs):
        """Registers a blueprint."""
        current_app.register_blueprint(blueprint, **kwargs)

    def create_table(self, model, db):
        """Creates the relation for the model

        :param model: The Model which should be created
        :param db: The database instance.
        """
        if not model.__table__.exists(bind=db.engine):
            model.__table__.create(bind=db.engine)

    def drop_table(self, model, db):
        """Drops the relation for the bounded model.

        :param model: The model on which the table is bound.
        :param db: The database instance.
        """
        model.__table__.drop(bind=db.engine)

    def create_all_tables(self, models, db):
        """A interface for creating all models specified in ``models``.

        :param models: A list with models
        :param db: The database instance
        """
        for model in models:
            self.create_table(model, db)

    def drop_all_tables(self, models, db):
        """A interface for dropping all models specified in the
        variable ``models``.

        :param models: A list with models
        :param db: The database instance.
        """
        for model in models:
            self.drop_table(model, db)
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import

from math import ceil
from operator import itemgetter

from flask_restplus import inputs

from flexget.api import api, APIResource, jsonify, request
from flexget.plugins.filter import seen
from flexget.utils.imdb import is_imdb_url, extract_id

def return_imdb_id(value):
    if is_imdb_url(value):
        imdb_id = extract_id(value)
        if imdb_id:
            value = imdb_id
    return value

seen_api = api.namespace('seen', description='Managed Flexget seen entries and fields')

seen_field_object = {
    'type': 'object',
    'properties': {
        'id': {'type': 'integer'},
        'field': {'type': 'string'},
        'value': {'type': 'string'},
        'added': {'type': 'string'},
        'seen_entry_id': {'type': 'integer'}
    }
}

seen_object = {
    'type': 'object',
    'properties': {
        'id': {'type': 'integer'},
        'title': {'type': 'string'},
        'reason': {'type': 'string'},
        'task': {'type': 'string'},
        'added': {'type': 'string'},
        'local': {'type': 'string'},
        'fields': {'type': 'array', 'items': seen_field_object}
    }
}
seen_object_schema = api.schema('seen_object_schema', seen_object)

seen_object_input_schema = {
    'type': 'object',
    'properties': {
        'title': {'type': 'string'},
        'reason': {'type': 'string'},
        'task': {'type': 'string'},
        'local': {'type': 'boolean', 'default': False},
        'fields': {'type': 'object'}
    },
    'required': ['title', 'fields', 'task'],
    'additional_properties': False
}
seen_object_input_schema = api.schema('seen_object_input_schema', seen_object_input_schema)

seen_search_schema = {
    'type': 'object',
    'properties': {
        'seen_entries': {
            'type': 'array',
            'items': seen_object
        },
        'number_of_seen_entries': {'type': 'integer'},
        'total_number_of_pages': {'type': 'integer'},
        'page_number': {'type': 'integer'}
    }
}
seen_search_schema = api.schema('seen_search_schema', seen_search_schema)

def seen_search_local_status_enum(value):
    try:
        return inputs.boolean(value)
    except ValueError:
        if value != 'all':
            raise ValueError('Invalid value received')
        return value

seen_search_sort_enum_list = ['title', 'task', 'added', 'local', 'id']

def seen_search_sort_enum(value):
    enum = seen_search_sort_enum_list
    if value not in enum:
        raise ValueError('Value expected to be in' + ' ,'.join(enum))
    return value

def seen_search_sort_order_enum(value):
    enum = ['desc', 'asc']
    if isinstance(value, bool):
        return value
    if value not in enum:
        raise ValueError('Value expected to be in' + ' ,'.join(enum))
    if value == 'desc':
        return True
    return False

seen_search_parser = api.parser()
seen_search_parser.add_argument('page', type=int, default=1, help='Page number')
seen_search_parser.add_argument('max', type=int, default=100, help='SEen entries per page')
seen_search_parser.add_argument('local_seen', type=seen_search_local_status_enum, default='all',
                                help='Filter list by local status. Filter by true, false or all. Default is all')
seen_search_parser.add_argument('sort_by', type=seen_search_sort_enum, default='added',
                                help="Sort response by {0}".format(' ,'.join(seen_search_sort_enum_list)))
seen_search_parser.add_argument('order', type=seen_search_sort_order_enum, default='desc',
                                help='Sorting order. Can be asc or desc. Default is desc')

@seen_api.route('/<string:value>')
@api.doc(params={'value': 'Name, IMDB ID or IMDB URL'})
class SeenSEarchAPI(APIResource):
    @api.response(404, 'Page does not exist')
    @api.response(200, 'Successfully retrieved seen objects', seen_search_schema)
    @api.doc(parser=seen_search_parser)
    def get(self, value, session):
        """ Search for seen entries """
        args = seen_search_parser.parse_args()
        page = args['page']
        max_results = args['max']
        status = args['local_seen']
        sort_by = args['sort_by']
        order = args['order']
        # Handles default if it explicitly called
        if order == 'desc':
            order = True

        value = return_imdb_id(value)
        value = '%' + value + '%'
        seen_entries_list = seen.search(value, status, session)
        count = len(seen_entries_list)

        pages = int(ceil(count / float(max_results)))
        seen_entries = []
        if page > pages and pages != 0:
            return {'error': 'page %s does not exist' % page}, 404

        start = (page - 1) * max_results
        finish = start + max_results
        if finish > count:
            finish = count

        for seen_entry_num in range(start, finish):
            seen_entries.append(seen_entries_list[seen_entry_num].to_dict())

        sorted_seen_entries_list = sorted(seen_entries, key=itemgetter(sort_by), reverse=order)

        return jsonify({
            'seen_entries': sorted_seen_entries_list,
            'number_of_seen_entries': count,
            'page_number': page,
            'total_number_of_pages': pages
        })

@seen_api.route('/')
class SeenAddAPI(APIResource):
    @api.response(400, 'A matching seen object is already added')
    @api.response(200, 'Successfully added new seen object', seen_object_schema)
    @api.validate(seen_object_input_schema)
    def post(self, session):
        """ Manually add entries to seen plugin """
        data = request.json
        kwargs = {
            'title': data.get('title'),
            'task_name': 'seen_API',
            'fields': data.get('fields'),
            'reason': data.get('reason'),
            'local': data.get('local', False),
            'session': session
        }
        # TODO check if seen entry already exist
        seen_entry = seen.add(**kwargs)
        return jsonify({
            'status': 'success',
            'message': 'successfully added seen object',
            'seen_object': seen_entry
        })
<commit_msg>Added check for existing value<commit_after>from __future__ import unicode_literals, division, absolute_import

from math import ceil
from operator import itemgetter

from flask_restplus import inputs

from flexget.api import api, APIResource, jsonify, request
from flexget.plugins.filter import seen
from flexget.utils.imdb import is_imdb_url, extract_id

def return_imdb_id(value):
    if is_imdb_url(value):
        imdb_id = extract_id(value)
        if imdb_id:
            value = imdb_id
    return value

seen_api = api.namespace('seen', description='Managed Flexget seen entries and fields')

seen_field_object = {
    'type': 'object',
    'properties': {
        'id': {'type': 'integer'},
        'field': {'type': 'string'},
        'value': {'type': 'string'},
        'added': {'type': 'string'},
        'seen_entry_id': {'type': 'integer'}
    }
}

seen_object = {
    'type': 'object',
    'properties': {
        'id': {'type': 'integer'},
        'title': {'type': 'string'},
        'reason': {'type': 'string'},
        'task': {'type': 'string'},
        'added': {'type': 'string'},
        'local': {'type': 'string'},
        'fields': {'type': 'array', 'items': seen_field_object}
    }
}
seen_object_schema = api.schema('seen_object_schema', seen_object)

seen_object_input_schema = {
    'type': 'object',
    'properties': {
        'title': {'type': 'string'},
        'reason': {'type': 'string'},
        'task': {'type': 'string'},
        'local': {'type': 'boolean', 'default': False},
        'fields': {'type': 'object'}
    },
    'required': ['title', 'fields', 'task'],
    'additional_properties': False
}
seen_object_input_schema = api.schema('seen_object_input_schema', seen_object_input_schema)

seen_search_schema = {
    'type': 'object',
    'properties': {
        'seen_entries': {
            'type': 'array',
            'items': seen_object
        },
        'number_of_seen_entries': {'type': 'integer'},
        'total_number_of_pages': {'type': 'integer'},
        'page_number': {'type': 'integer'}
    }
}
seen_search_schema = api.schema('seen_search_schema', seen_search_schema)

def seen_search_local_status_enum(value):
    try:
        return inputs.boolean(value)
    except ValueError:
        if value != 'all':
            raise ValueError('Invalid value received')
        return value

seen_search_sort_enum_list = ['title', 'task', 'added', 'local', 'id']

def seen_search_sort_enum(value):
    enum = seen_search_sort_enum_list
    if value not in enum:
        raise ValueError('Value expected to be in' + ' ,'.join(enum))
    return value

def seen_search_sort_order_enum(value):
    enum = ['desc', 'asc']
    if isinstance(value, bool):
        return value
    if value not in enum:
        raise ValueError('Value expected to be in' + ' ,'.join(enum))
    if value == 'desc':
        return True
    return False

seen_search_parser = api.parser()
seen_search_parser.add_argument('page', type=int, default=1, help='Page number')
seen_search_parser.add_argument('max', type=int, default=100, help='SEen entries per page')
seen_search_parser.add_argument('local_seen', type=seen_search_local_status_enum, default='all',
                                help='Filter list by local status. Filter by true, false or all. Default is all')
seen_search_parser.add_argument('sort_by', type=seen_search_sort_enum, default='added',
                                help="Sort response by {0}".format(' ,'.join(seen_search_sort_enum_list)))
seen_search_parser.add_argument('order', type=seen_search_sort_order_enum, default='desc',
                                help='Sorting order. Can be asc or desc. Default is desc')

@seen_api.route('/<string:value>')
@api.doc(params={'value': 'Name, IMDB ID or IMDB URL'})
class SeenSEarchAPI(APIResource):
    @api.response(404, 'Page does not exist')
    @api.response(200, 'Successfully retrieved seen objects', seen_search_schema)
    @api.doc(parser=seen_search_parser)
    def get(self, value, session):
        """ Search for seen entries """
        args = seen_search_parser.parse_args()
        page = args['page']
        max_results = args['max']
        status = args['local_seen']
        sort_by = args['sort_by']
        order = args['order']
        # Handles default if it explicitly called
        if order == 'desc':
            order = True

        value = return_imdb_id(value)
        value = '%' + value + '%'
        seen_entries_list = seen.search(value, status, session)
        count = len(seen_entries_list)

        pages = int(ceil(count / float(max_results)))
        seen_entries = []
        if page > pages and pages != 0:
            return {'error': 'page %s does not exist' % page}, 404

        start = (page - 1) * max_results
        finish = start + max_results
        if finish > count:
            finish = count

        for seen_entry_num in range(start, finish):
            seen_entries.append(seen_entries_list[seen_entry_num].to_dict())

        sorted_seen_entries_list = sorted(seen_entries, key=itemgetter(sort_by), reverse=order)

        return jsonify({
            'seen_entries': sorted_seen_entries_list,
            'number_of_seen_entries': count,
            'page_number': page,
            'total_number_of_pages': pages
        })

@seen_api.route('/')
class SeenAddAPI(APIResource):
    @api.response(400, 'A matching seen object is already added')
    @api.response(200, 'Successfully added new seen object', seen_object_schema)
    @api.validate(seen_object_input_schema)
    def post(self, session):
        """ Manually add entries to seen plugin """
        data = request.json
        kwargs = {
            'title': data.get('title'),
            'task_name': 'seen_API',
            'fields': data.get('fields'),
            'reason': data.get('reason'),
            'local': data.get('local', False),
            'session': session
        }
        values = [value for value in kwargs['fields'].values()]
        exist = seen.search_by_field_values(field_value_list=values, task_name='seen_API', local=kwargs['local'],
                                            session=session)
        if exist:
            return {'status': 'error',
                    'message': "Seen entry matching the value '{0}' is already added".format(exist.value)}, 400

        seen_entry = seen.add(**kwargs)
        return jsonify({
            'status': 'success',
            'message': 'successfully added seen object',
            'seen_object': seen_entry
        })
<|endoftext|>