<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
    https://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.

Writer file used to serialize and save data to TFRecords files.
"""
import numpy as np
import os
import pandas as pd
import tensorflow as tf

def add_candidate_content(context, candidates):
  """Add the tensor for the classes this particular video is a candidate for.

  Args:
    context: context of the video
    candidates: dictionary of candidates. Key is video id and value is list of candidate classes
  """
  video_id = tf.convert_to_tensor(context["id"])[0].numpy()
  if video_id in candidates.keys():
    context["candidate_labels"] = tf.convert_to_tensor(candidates[video_id])
  else:
    context["candidate_labels"] = tf.convert_to_tensor([])
  return context

def convert_labels(labels, class_csv="vocabulary.csv"):
  """Convert labels from range [0,3861] to range [0,1000]

  Args:
    labels: Tensor of labels to be converted
    class_csv: csv file containing conversion details
  """
  class_dataframe = pd.read_csv(class_csv, index_col=0)
  class_indices = class_dataframe.index.tolist()
  class_indices = np.array(class_indices)

  labels = labels.numpy()
  new_labels = []
  for label in labels:
    check = np.nonzero(class_indices == label)[0]
    if len(check) > 0:
      new_label = check.tolist()[0]
      new_labels.append(new_label)
  return tf.convert_to_tensor(new_labels)

def convert_to_feature(item, type):
  """Convert item to FeatureList.

  Args:
    item: item to be converted
    type: string denoting the type of item. Can be "float", "byte", or "int"
  """
  if type == "float":
    item = tf.train.FloatList(value=item)
    item = tf.train.Feature(float_list=item)
  elif type == "byte":
    item = tf.train.BytesList(value=item)
    item = tf.train.Feature(bytes_list=item)
  elif type == "int":
    item = tf.train.Int64List(value=item)
    item = tf.train.Feature(int64_list=item)
  else:
    print("Invalid type entered for converting feature")
  return item

def serialize_features(features):
  """Serialize features.

  Args:
    features: features of the video
  """
  audio = features["audio"][0].numpy().tostring()
  rgb = features["rgb"][0].numpy().tostring()
  audio = convert_to_feature([audio], "byte")
  rgb = convert_to_feature([rgb], "byte")
  features = {"audio": tf.train.FeatureList(feature=[audio]), "rgb": tf.train.FeatureList(feature=[rgb])}
  features = tf.train.FeatureLists(feature_list=features)
  return features

def serialize_video_context(context):
  """Serialize context for a video.

  Args:
    context: context of the video
  """
  video_id = tf.convert_to_tensor(context["id"])[0]
  labels = context["labels"].values
  segment_labels = context["segment_labels"].values
  segment_start_times = context["segment_start_times"].values
  segment_scores = context["segment_scores"].values
  candidate_labels = context["candidate_labels"]
  labels =  convert_labels(labels)
  segment_labels = convert_labels(segment_labels)

  context["id"] = convert_to_feature([video_id.numpy()], "byte")
  context["labels"] = convert_to_feature(labels.numpy(), "int")
  context["segment_labels"] = convert_to_feature(segment_labels.numpy(), "int")
  context["segment_start_times"] = convert_to_feature(segment_start_times.numpy(), "int")
  context["segment_score"] = convert_to_feature(segment_scores.numpy(), "float")

  context = tf.train.Features(feature=context)
  return context

def serialize_segment_context(context):
  """Serialize context for a segment.

  Args:
    context: context of the video
  """
  video_id = tf.convert_to_tensor(context["id"])[0]
  segment_label = context["segment_label"]
  segment_start_time = context["segment_start_time"]
  segment_score = context["segment_score"]
  segment_label = convert_labels(segment_label)

  context["id"] = convert_to_feature([video_id.numpy()], "byte")
  context["segment_label"] = convert_to_feature(segment_label.numpy(), "int")
  context["segment_start_time"] = convert_to_feature(segment_start_time.numpy(), "int")
  context["segment_score"] = convert_to_feature(segment_score.numpy(), "float")
  context = tf.train.Features(feature=context)
  return context

def serialize_data(context, features, type):
  """Serialize video or segment from context and features.

  Args:
    context: context of the video
    features: features of the video
    type: type of data to store. Can either be video or segment
  """
  features = serialize_features(features)
  if type == "video":
    context = serialize_video_context(context)
  elif type == "segment":
    context = serialize_segment_context(context)
  else:
    print("Incorrect type chosen for serialization.")
  example = tf.train.SequenceExample(feature_lists=features, context=context)
  return example.SerializeToString()

def save_shard(data_dir, shard, file_type, shard_number):
  """Save a shard data to data_dir as a TFRecords file.

  Args:
    data_dir: directory for file to be saved
    shard: list of serialized examples to be saved
    file_type: prefix of file name.
    shard_number: suffix of file name.
  """
  print(f"Processing shard number {shard_number}")
  shard = tf.convert_to_tensor(shard)
  shard_dataset = tf.data.Dataset.from_tensor_slices(shard)
  file_name = file_type + str(shard_number)
  file_path = os.path.join(data_dir, '%s.tfrecord' % file_name)
  writer = tf.data.experimental.TFRecordWriter(file_path)
  writer.write(shard_dataset)

def save_data(new_data_dir, input_dataset, candidates, file_type="test", shard_size=17):
  """Save data as TFRecords Datasets in new_data_dir.

  Args:
    new_data_dir: string giving the directory to save TFRecords Files
    input_dataset: original dataset before candidate generation
    candidates: list of lists where each inner list contains the class indices that the corresponding input data is a candidate for. len(candidates) == len(input_dataset)
  """
  shard_counter = 0
  shard_number = 0
  shard = []
  for video in input_dataset:
    context = video[0]
    features = video[1]
    context = add_candidate_content(context, candidates)
    serialized_video = serialize_data(context, features, "video")
    shard.append(serialized_video)
    shard_counter += 1
    if shard_counter == shard_size:
      save_shard(new_data_dir, shard, file_type, shard_number)
      shard_counter = 0
      shard_number += 1
      shard = []
  #Handles overflow
  if shard_counter != 0:
    save_shard(new_data_dir, shard, file_type, shard_number)
    shard_counter = 0
    shard_number += 1
    shard = []

def split_data(data_dir, input_dataset, shard_size=85, num_classes=1000, file_type="class"):
  """Save data as TFRecords Datasets in new_data_dir.

  Args:
    new_data_dir: string giving the directory to save TFRecords Files
    input_dataset: original dataset before candidate generation
  """
  #Context: id, label, score
  #Features: rgb, audio for 1 segment
  #Convert input_dataset from video level data to multiple segments.
  video_holder = [[]]*num_classes
  print(len(video_holder))
  video_number = 0
  for video in input_dataset:
    print(f"Processing video number {video_number}")
    context = video[0]
    features = video[1]
    segment_start_times = context["segment_start_times"].values.numpy()
    for segment_index in range(len(segment_start_times)):
      new_context = {}
      new_context["id"] = context["id"]
      new_context["segment_label"] = tf.convert_to_tensor([context["segment_labels"].values.numpy()[segment_index]])
      new_context["segment_start_time"] = tf.convert_to_tensor([segment_start_times[segment_index]])
      new_context["segment_score"] = tf.convert_to_tensor([context["segment_scores"].values.numpy()[segment_index]])
      new_features = {}
      new_features["rgb"] = features["rgb"][:,segment_start_times[segment_index]:segment_start_times[segment_index]+5,:]
      new_features["audio"] = features["audio"][:,segment_start_times[segment_index]:segment_start_times[segment_index]+5,:]

      label = new_context["segment_label"]
      label = convert_labels(label).numpy()[0]
      serialized_video = serialize_data(new_context, new_features, "segment")
      video_holder[label].append(serialized_video)
      print(len(video_holder[201]))
      print(label)
    video_number += 1
  print(len(video_holder[200]))
  assert False
  for shard_number in range(len(video_holder)):
    save_shard(data_dir, video_holder[shard_number], file_type, shard_number)<commit_msg>Class Specific Features<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
    https://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.

Writer file used to serialize and save data to TFRecords files.
"""
import numpy as np
import os
import pandas as pd
import tensorflow as tf

def add_candidate_content(context, candidates):
  """Add the tensor for the classes this particular video is a candidate for.

  Args:
    context: context of the video
    candidates: dictionary of candidates. Key is video id and value is list of candidate classes
  """
  video_id = tf.convert_to_tensor(context["id"])[0].numpy()
  if video_id in candidates.keys():
    context["candidate_labels"] = tf.convert_to_tensor(candidates[video_id])
  else:
    context["candidate_labels"] = tf.convert_to_tensor([])
  return context

def convert_labels(labels, class_csv="vocabulary.csv"):
  """Convert labels from range [0,3861] to range [0,1000]

  Args:
    labels: Tensor of labels to be converted
    class_csv: csv file containing conversion details
  """
  class_dataframe = pd.read_csv(class_csv, index_col=0)
  class_indices = class_dataframe.index.tolist()
  class_indices = np.array(class_indices)

  labels = labels.numpy()
  new_labels = []
  for label in labels:
    check = np.nonzero(class_indices == label)[0]
    if len(check) > 0:
      new_label = check.tolist()[0]
      new_labels.append(new_label)
  return tf.convert_to_tensor(new_labels)

def convert_to_feature(item, type):
  """Convert item to FeatureList.

  Args:
    item: item to be converted
    type: string denoting the type of item. Can be "float", "byte", or "int"
  """
  if type == "float":
    item = tf.train.FloatList(value=item)
    item = tf.train.Feature(float_list=item)
  elif type == "byte":
    item = tf.train.BytesList(value=item)
    item = tf.train.Feature(bytes_list=item)
  elif type == "int":
    item = tf.train.Int64List(value=item)
    item = tf.train.Feature(int64_list=item)
  else:
    print("Invalid type entered for converting feature")
  return item

def serialize_features(features):
  """Serialize features.

  Args:
    features: features of the video
  """
  audio = features["audio"][0].numpy().tostring()
  rgb = features["rgb"][0].numpy().tostring()
  audio = convert_to_feature([audio], "byte")
  rgb = convert_to_feature([rgb], "byte")
  features = {"audio": tf.train.FeatureList(feature=[audio]), "rgb": tf.train.FeatureList(feature=[rgb])}
  features = tf.train.FeatureLists(feature_list=features)
  return features

def serialize_video_context(context):
  """Serialize context for a video.

  Args:
    context: context of the video
  """
  video_id = tf.convert_to_tensor(context["id"])[0]
  labels = context["labels"].values
  segment_labels = context["segment_labels"].values
  segment_start_times = context["segment_start_times"].values
  segment_scores = context["segment_scores"].values
  candidate_labels = context["candidate_labels"]
  labels =  convert_labels(labels)
  segment_labels = convert_labels(segment_labels)

  context["id"] = convert_to_feature([video_id.numpy()], "byte")
  context["labels"] = convert_to_feature(labels.numpy(), "int")
  context["segment_labels"] = convert_to_feature(segment_labels.numpy(), "int")
  context["segment_start_times"] = convert_to_feature(segment_start_times.numpy(), "int")
  context["segment_score"] = convert_to_feature(segment_scores.numpy(), "float")

  context = tf.train.Features(feature=context)
  return context

def serialize_segment_context(context):
  """Serialize context for a segment.

  Args:
    context: context of the video
  """
  video_id = tf.convert_to_tensor(context["id"])[0]
  segment_label = context["segment_label"]
  segment_start_time = context["segment_start_time"]
  segment_score = context["segment_score"]
  segment_label = convert_labels(segment_label)

  context["id"] = convert_to_feature([video_id.numpy()], "byte")
  context["segment_label"] = convert_to_feature(segment_label.numpy(), "int")
  context["segment_start_time"] = convert_to_feature(segment_start_time.numpy(), "int")
  context["segment_score"] = convert_to_feature(segment_score.numpy(), "float")
  context = tf.train.Features(feature=context)
  return context

def serialize_data(context, features, type):
  """Serialize video or segment from context and features.

  Args:
    context: context of the video
    features: features of the video
    type: type of data to store. Can either be video or segment
  """
  features = serialize_features(features)
  if type == "video":
    context = serialize_video_context(context)
  elif type == "segment":
    context = serialize_segment_context(context)
  else:
    print("Incorrect type chosen for serialization.")
  example = tf.train.SequenceExample(feature_lists=features, context=context)
  return example.SerializeToString()

def save_shard(data_dir, shard, file_type, shard_number):
  """Save a shard data to data_dir as a TFRecords file.

  Args:
    data_dir: directory for file to be saved
    shard: list of serialized examples to be saved
    file_type: prefix of file name.
    shard_number: suffix of file name.
  """
  print(f"Processing shard number {shard_number}")
  shard = tf.convert_to_tensor(shard)
  shard_dataset = tf.data.Dataset.from_tensor_slices(shard)
  file_name = file_type + str(shard_number)
  file_path = os.path.join(data_dir, '%s.tfrecord' % file_name)
  writer = tf.data.experimental.TFRecordWriter(file_path)
  writer.write(shard_dataset)

def save_data(new_data_dir, input_dataset, candidates, file_type="test", shard_size=17):
  """Save data as TFRecords Datasets in new_data_dir.

  Args:
    new_data_dir: string giving the directory to save TFRecords Files
    input_dataset: original dataset before candidate generation
    candidates: list of lists where each inner list contains the class indices that the corresponding input data is a candidate for. len(candidates) == len(input_dataset)
  """
  shard_counter = 0
  shard_number = 0
  shard = []
  for video in input_dataset:
    context = video[0]
    features = video[1]
    context = add_candidate_content(context, candidates)
    serialized_video = serialize_data(context, features, "video")
    shard.append(serialized_video)
    shard_counter += 1
    if shard_counter == shard_size:
      save_shard(new_data_dir, shard, file_type, shard_number)
      shard_counter = 0
      shard_number += 1
      shard = []
  #Handles overflow
  if shard_counter != 0:
    save_shard(new_data_dir, shard, file_type, shard_number)
    shard_counter = 0
    shard_number += 1
    shard = []

def split_data(data_dir, input_dataset, shard_size=85, num_classes=1000, file_type="class"):
  """Save data as TFRecords Datasets in new_data_dir.

  Args:
    new_data_dir: string giving the directory to save TFRecords Files
    input_dataset: original dataset before candidate generation
  """
  #Context: id, label, score
  #Features: rgb, audio for 1 segment
  #Convert input_dataset from video level data to multiple segments.
  video_holder = [[]]*num_classes
  print(len(video_holder))
  video_number = 0
  for video in input_dataset:
    print(f"Processing video number {video_number}")
    context = video[0]
    features = video[1]
    segment_start_times = context["segment_start_times"].values.numpy()
    for segment_index in range(len(segment_start_times)):
      new_context = {}
      new_context["id"] = context["id"]
      new_context["segment_label"] = tf.convert_to_tensor([context["segment_labels"].values.numpy()[segment_index]])
      new_context["segment_start_time"] = tf.convert_to_tensor([segment_start_times[segment_index]])
      new_context["segment_score"] = tf.convert_to_tensor([context["segment_scores"].values.numpy()[segment_index]])
      new_features = {}
      new_features["rgb"] = features["rgb"][:,segment_start_times[segment_index]:segment_start_times[segment_index]+5,:]
      new_features["audio"] = features["audio"][:,segment_start_times[segment_index]:segment_start_times[segment_index]+5,:]

      label = new_context["segment_label"]
      label = convert_labels(label).numpy()[0]
      serialized_video = serialize_data(new_context, new_features, "segment")
      video_holder[label].append(serialized_video)
      print(label)
    video_number += 1
  for i in range(len(video_holder)):
    print(len(video_holder[i]))
  print(len(video_holder))
  assert False
  for shard_number in range(len(video_holder)):
    save_shard(data_dir, video_holder[shard_number], file_type, shard_number)<|endoftext|>


<commit_before>"""
This module provides an interface to the mechanisms
used to resolve module name from mofule filepath.

    :copyright: 2008 by Takanori Ishikawa
    :license: MIT, see LICENSE for more details.
"""

import os
import sys
from os.path import isdir, abspath, expanduser
from modipyd import utils

def normalize_path(filepath):
    """Normalize path name"""
    #return os.path.realpath(abspath(expanduser(filepath)))
    return abspath(expanduser(filepath))

def resolve_relative_modulename(modulename, package, level):
    if level <= 0:
        return modulename
    else:
        for _ in xrange(1, level):
            package = package[:package.rindex('.')]

        return '.'.join((package, modulename))

class ModuleNameResolver(object):
    """
    This class provides an interface to the mechanisms
    used to resolve module name from a python source file path.
    """

    def __init__(self, search_paths=None):
        """
        The ``search_paths`` argument is module search path,
        if *search_paths* is omitted or ``None``, ``sys.path`` is used.
        """
        super(ModuleNameResolver, self).__init__()

        # Cofigure module search path (copy it)
        syspaths = (search_paths or sys.path)
        syspaths = utils.wrap_sequence(syspaths)
        assert isinstance(syspaths, (list, tuple))

        self.search_paths = [normalize_path(d) for d in syspaths if isdir(d)]
        for d in self.search_paths:
            assert isinstance(d, basestring)
            assert isdir(d)

        # caches
        self.cache_package = {}

    def python_package(self, directory):
        try:
            return self.cache_package[directory]
        except KeyError:
            package = utils.python_package(directory)
            self.cache_package[directory] = package
            return package

    def resolve(self, filepath):
        """
        Resolve the module name from *filepath* on search_paths.
        Return (module name, package name), if the module is not
        in a package, package name is ``None``.
        """

        if not filepath:
            raise ImportError("filepath is empty")

        filepath = normalize_path(filepath)
        if not utils.python_module_file(filepath):
            raise ImportError("Not a python script: %s" % filepath)

        # Searching...
        #
        # split path components
        # ignore extention (e.g. '.py', '.pyc')
        dirpath, modname = os.path.split(filepath)
        modname, _ = os.path.splitext(modname)

        skipped_name = None
        package_name = None
        compiled_forms = [(dirpath, [modname])]

        for syspath in self.search_paths:

            i = 0
            while not compiled_forms[i][0] == syspath:
                # compile
                i += 1
                if i == len(compiled_forms):
                    d, names = compiled_forms[i-1]

                    # Encountered not a package, or not found in
                    # the module search path
                    if not self.python_package(d) or d == '/':
                        break

                    names = names[:] # copy
                    d, parent = os.path.split(d)
                    names.insert(0, parent)

                    compiled_forms.append((d, names))
            else:
                d, names = compiled_forms[i]
                if self.python_package(dirpath):
                    package_name = '.'.join(names[:-1])
                else:
                    package_name = None

                if i == 0:
                    if self.python_package(d):
                        # script is created under a package,
                        # but its module name is not a package style
                        # (e.g. 'package.module').
                        #
                        # For example, consider:
                        #
                        #   src/packageA/__init__.py
                        #               /a.py
                        #
                        # and *sys.path* is:
                        #
                        #   ['src/packageA', 'src']
                        #
                        # a.py is in src/packageA, so module name 'a' is
                        # matched. But better module name is 'packageA.a'
                        # because a.py is created under the package
                        # 'packageA'.
                        #
                        # Store result name, and continues with
                        # the next iteration of the **for** loop.
                        skipped_name = names[0]
                        package_name = None
                        continue
                else:
                    if names[-1] == '__init__':
                        # The package initialization module.
                        # Remove tail '__init__'
                        del names[-1]

                return '.'.join(names), package_name

        if skipped_name is not None:
            return skipped_name, package_name
        else:
            raise ImportError("No module name found: %s" % filepath)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Addressed module/script name collision problem: validates resolved module name with imp.find_module (ticket #8 at http://metareal.lighthouseapp.com/projects/17658-modipyd/tickets/8)<commit_after>"""
This module provides an interface to the mechanisms
used to resolve module name from module filepath.

    :copyright: 2008 by Takanori Ishikawa
    :license: MIT, see LICENSE for more details.
"""

import os
import sys
import imp
from os.path import isdir, abspath, expanduser
from modipyd import utils

def normalize_path(filepath):
    """Normalize path name"""
    return abspath(expanduser(filepath))

def resolve_relative_modulename(modulename, package, level):
    if level <= 0:
        return modulename
    else:
        for _ in xrange(1, level):
            package = package[:package.rindex('.')]

        return '.'.join((package, modulename))

class ModuleNameResolver(object):
    """
    This class provides an interface to the mechanisms
    used to resolve module name from a python source file path.
    """

    def __init__(self, search_paths=None):
        """
        The ``search_paths`` argument is module search path,
        if *search_paths* is omitted or ``None``, ``sys.path`` is used.
        """
        super(ModuleNameResolver, self).__init__()

        # Cofigure module search path (copy it)
        syspaths = (search_paths or sys.path)
        syspaths = utils.wrap_sequence(syspaths)
        assert isinstance(syspaths, (list, tuple))

        self.search_paths = [normalize_path(d) for d in syspaths if isdir(d)]
        for d in self.search_paths:
            assert isinstance(d, basestring)
            assert isdir(d)

        # caches
        self.cache_package = {}

    def python_package(self, directory):
        try:
            return self.cache_package[directory]
        except KeyError:
            package = utils.python_package(directory)
            self.cache_package[directory] = package
            return package

    def resolve(self, filepath):
        """
        Resolves the name of the module located at *filepath*.
        If successful, returns names of module and package.
        """
        modname, package = self._resolve(filepath)
        if modname:
            # validates resolved module name with imp.find_module
            module_path = list(self.search_paths)
            description = None
            for name in modname.split('.'):
                file, pathname, description = imp.find_module(name, module_path)
                module_path = [ pathname ]
                if file:
                    file.close()
                if description[2] != imp.PKG_DIRECTORY:
                    break

            module_path = module_path[0]
            assert description
            if description[2] == imp.PKG_DIRECTORY:
                module_path = os.path.join(module_path, '__init__.py')

            pt1, _ = os.path.splitext(normalize_path(filepath))
            pt2, _ = os.path.splitext(normalize_path(module_path))

            if pt1 != pt2:
                raise ImportError("Can't resolve module name: %s" % filepath)

        return modname, package

    def _resolve(self, filepath):

        if not filepath:
            raise ImportError("filepath is empty")

        filepath = normalize_path(filepath)
        if not utils.python_module_file(filepath):
            raise ImportError("Not a python script: %s" % filepath)

        # Searching...
        #
        # split path components
        # ignore extention (e.g. '.py', '.pyc')
        dirpath, modname = os.path.split(filepath)
        modname, _ = os.path.splitext(modname)

        skipped_name = None
        package_name = None
        compiled_forms = [(dirpath, [modname])]

        for syspath in self.search_paths:

            i = 0
            while not compiled_forms[i][0] == syspath:
                # compile
                i += 1
                if i == len(compiled_forms):
                    d, names = compiled_forms[i-1]

                    # Encountered not a package, or not found in
                    # the module search path
                    if not self.python_package(d) or d == '/':
                        break

                    names = names[:] # copy
                    d, parent = os.path.split(d)
                    names.insert(0, parent)

                    compiled_forms.append((d, names))
            else:
                d, names = compiled_forms[i]
                if self.python_package(dirpath):
                    package_name = '.'.join(names[:-1])
                else:
                    package_name = None

                if i == 0:
                    if self.python_package(d):
                        # script is created under a package,
                        # but its module name is not a package style
                        # (e.g. 'package.module').
                        #
                        # For example, consider:
                        #
                        #   src/packageA/__init__.py
                        #               /a.py
                        #
                        # and *sys.path* is:
                        #
                        #   ['src/packageA', 'src']
                        #
                        # a.py is in src/packageA, so module name 'a' is
                        # matched. But better module name is 'packageA.a'
                        # because a.py is created under the package
                        # 'packageA'.
                        #
                        # Store result name, and continues with
                        # the next iteration of the **for** loop.
                        skipped_name = names[0]
                        package_name = None
                        continue
                else:
                    if names[-1] == '__init__':
                        # The package initialization module.
                        # Remove tail '__init__'
                        del names[-1]

                return '.'.join(names), package_name

        if skipped_name is not None:
            return skipped_name, package_name
        else:
            raise ImportError("No module name found: %s" % filepath)

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>


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

###########################################################
# Copyright 2011 Daniel 'grindhold' Brendle and Team
#
# This file is part of Scoville.
#
# Scoville is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later
# version.
#
# Scoville is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public
# License along with Scoville.
# If not, see http://www.gnu.org/licenses/.
###########################################################

from hashlib import sha256
from datetime import datetime, timedelta
from random import randrange

class SessionException(Exception):
    ERRORS = {
        0:"""SessionError: Session Expired""",
        1:"""SessionError: Can only attach user to session""",
        2:"""SessionError: This Session does not exist"""
    }

    @classmethod
    def get_msg(cls,nr, info=""):
        return "SE_"+str(nr)+": "+cls.ERRORS[nr]+" "+info

class SessionManager(object):
    """
    Wraps Session's classmethod
    """
    def __init__(self, core):
        self._core = core
        Session.set_core(core)

        self.get_session = Session.get_session
        self.create_session = Session.create_session
        self.get_current_session = Session.get_current_session
        self.set_current_session = Session.set_current_session
        self.get_current_session_user = Session.get_current_session_user

class Session(object):
    """
    Handles clientsessions by a given session_id
    TODO: Session saves environmental variablees
    TODO: Session takes env-vars in account when verifying the id of the client
    """
    ACTIVE_SESSIONS = {}
    CURRENT_SESSION = None
    @classmethod
    def set_core(cls,core):
        """
        trivial
        """
        cls._core = core

    @classmethod
    def get_session(cls,session_id):
        """
        returns the session if it's not expired or nonexistant
        """
        if not cls.ACTIVE_SESSIONS.has_key(session_id):
            raise SessionException(SessionException.get_msg(2))
        session = cls.ACTIVE_SESSIONS[session_id]
        if session._expiration < datetime.now():
            del(cls.ACTIVE_SESSIONS[session_id])
            raise SessionException(SessionException.get_msg(0))
        else:
            return session

    @classmethod
    def create_session(cls, user):
        """
        creates a session for a given user
        """
        return Session(cls._core,user)

    @classmethod
    def set_current_session(cls, session):
        """
        sets the session of the current user
        """
        cls.CURRENT_SESSION = session

    @classmethod
    def get_current_session(cls):
        """
        returns the current session of the current client
        """
        return cls.CURRENT_SESSION

    @classmethod
    def get_current_session_user(cls):
        """
        returns the user of the current clientsession
        """
        return cls.CURRENT_SESSION.get_user()

    def _generateRandomString(self,length):
        """
        generate  a random string
        TODO: outsource this function to some kind of helpers module (other occurence in database.py)
        """
        ret = ""
        for i in range(length):
            x = -1
            while x in (-1,91,92,93,94,95,96,60,61,62,58,59,63,64):
                x = randrange(48,123)
            ret+=chr(x)
        return ret

    def __init__(self,core, user):
        """
        initializes a session
        """
        self._core = core
        configuration = self._core.get_configuration() #TODO: Introduce Configvalue sessionduration

        self._id = sha256(self._generateRandomString(24)).hexdigest()
        self._user = user.get_id()
        self._expiration = datetime.now()+timedelta(0,int(configuration.get_entry("core.session_duration"))*3600)
        Session.ACTIVE_SESSIONS[self._id] = self

    def extend(self):
        """
        extends this session by a timespan defined as "core.session_extend" in hours in configuration
        """
        configuration = self._core.get_configuration()
        self._expiration += timedelta(0,int(configuration.get_entry("core.session_extend"))*3600)

    def get_user(self):
        """
        returns this session's user
        """
        usermanager = self._core.get_user_manager()
        return usermanager.get_user(self._user)

    def set_user(self,user):
        """
        set the user of this session
        """
        if type(user) == int:
            self._user = user
        elif user.__class__.__name__ == "User":
            self._user = user.get_id()
        else:
            raise SessionException(SessionException.get_msg(1))<commit_msg>sessions are stored in database<commit_after>#!/usr/bin/python
#-*- coding: utf-8 -*-

###########################################################
# Copyright 2011 Daniel 'grindhold' Brendle and Team
#
# This file is part of Scoville.
#
# Scoville is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later
# version.
#
# Scoville is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public
# License along with Scoville.
# If not, see http://www.gnu.org/licenses/.
###########################################################

from hashlib import sha256
from datetime import datetime, timedelta
from random import randrange

class SessionException(Exception):
    ERRORS = {
        0:"""SessionError: Session Expired""",
        1:"""SessionError: Can only attach user to session""",
        2:"""SessionError: This Session does not exist""",
        3:"""SessionError: Cannot Store a userless Session"""
    }

    @classmethod
    def get_msg(cls,nr, info=""):
        return "SE_"+str(nr)+": "+cls.ERRORS[nr]+" "+info

class SessionManager(object):
    """
    Wraps Session's classmethod
    """
    def __init__(self, core):
        self._core = core
        Session.set_core(core)

        self.get_session = Session.get_session
        self.create_session = Session.create_session
        self.get_current_session = Session.get_current_session
        self.set_current_session = Session.set_current_session
        self.get_current_session_user = Session.get_current_session_user

class Session(object):
    """
    Handles clientsessions by a given session_id
    TODO: Session saves environmental variablees
    TODO: Session takes env-vars in account when verifying the id of the client
    """
    CURRENT_SESSION = None
    @classmethod
    def set_core(cls,core):
        """
        trivial
        """
        cls._core = core

    @classmethod
    def get_session(cls,session_id):
        """
        returns the session if it's not expired or nonexistant
        """
        db = cls._core.get_db()
        stmnt = "SELECT SES_USR_ID, SES_EXPIRES FROM SESSIONS WHERE SES_ID = ? ;"

        cur = db.query(cls._core,stmnt,(session_id,))
        row = cur.fetchonemap()

        session=None

        if row is not None:
            user_manager = cls._core.get_user_manager()
            user = user_manager.get_user(row["SES_USR_ID"])
            session = Session(cĺs._core,user)
            session._id = session_id
            expiration = row["SES_EXPIRES"]
            if expiration < datetime.now():
                raise SessionException(SessionException.get_msg(0))
            session._expiration = row["SES_EXPIRES"]
        else:
            raise SessionException(SessionException.get_msg(2))
        return session

    @classmethod
    def create_session(cls, user):
        """
        creates a session for a given user
        """
        s = Session(cls._core,user)
        s.store()
        return s

    @classmethod
    def set_current_session(cls, session):
        """
        sets the session of the current user
        """
        cls.CURRENT_SESSION = session

    @classmethod
    def get_current_session(cls):
        """
        returns the current session of the current client
        """
        return cls.CURRENT_SESSION

    @classmethod
    def get_current_session_user(cls):
        """
        returns the user of the current clientsession
        """
        return cls.CURRENT_SESSION.get_user()

    def _generateRandomString(self,length):
        """
        generate  a random string
        TODO: outsource this function to some kind of helpers module (other occurence in database.py)
        """
        ret = ""
        for i in range(length):
            x = -1
            while x in (-1,91,92,93,94,95,96,60,61,62,58,59,63,64):
                x = randrange(48,123)
            ret+=chr(x)
        return ret

    def __init__(self,core, user):
        """
        initializes a session
        """
        self._core = core
        configuration = self._core.get_configuration() #TODO: Introduce Configvalue sessionduration

        self._id = sha256(self._generateRandomString(24)).hexdigest()
        self._user = user.get_id()
        self._expiration = datetime.now()+timedelta(0,int(configuration.get_entry("core.session_duration"))*3600)

    def extend(self):
        """
        extends this session by a timespan defined as "core.session_extend" in hours in configuration
        """
        configuration = self._core.get_configuration()
        self._expiration += timedelta(0,int(configuration.get_entry("core.session_extend"))*3600)

    def get_user(self):
        """
        returns this session's user
        """
        usermanager = self._core.get_user_manager()
        return usermanager.get_user(self._user)

    def set_user(self,user):
        """
        set the user of this session
        """
        if type(user) == int:
            self._user = user
        elif user.__class__.__name__ == "User":
            self._user = user.get_id()
        else:
            raise SessionException(SessionException.get_msg(1))

    def store(self):
        """
        stores this session in database
        """
        if self._user is None:
            raise SessionException(SessionException.get_msg(3))
        else:
            db = self._core.get_db()
            stmnt = "UPDATE OR INSERT INTO SESSIONS (SES_ID, SES_USR_ID, SES_EXPIRES) VALUES (?,?,?) MATCHING (SES_ID) ;"

            db.query(self._core,stmnt,(self._id,self._user.get_id(),self._expiration))

    def delete(self):
        """
        deletes this session
        """
        db = self._core.get_db()
        stmnt = "DELETE FROM SESSIONS WHERE SES_ID = ? ;"
        db.query(self._core,stmnt,(self._id,))
        <|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2011 The Pycroft Authors. See the AUTHORS file.
# This file is part of the Pycroft project and licensed under the terms of
# the Apache License, Version 2.0. See the LICENSE file for details.
"""
    pycroft.model.base
    ~~~~~~~~~~~~~~

    This module contains base stuff for all models.

    :copyright: (c) 2011 by AG DSN.
"""
from sqlalchemy import ForeignKey, MetaData
from sqlalchemy import Table, Column
from sqlalchemy.dialects import postgresql
from sqlalchemy.ext.declarative import DeclarativeMeta
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy.orm import relationship, backref
from sqlalchemy.types import BigInteger, Boolean, DateTime, Integer
from sqlalchemy.types import Text, String

_session = None

class _ModelMeta(DeclarativeMeta):
    """Metaclass for all mapped Database objects.
    """
    @property
    def q(cls):
        """This is a shortcut for easy querying of whole objects.

        With this metaclass shortcut you can query a Model with
        Model.q.filter(...) without using the verbose session stuff
        """
        global _session
        if _session is None:
            import pycroft.model.session
            _session = pycroft.model.session.session
        return _session.query(cls)

class _Base(object):
    """Baseclass for all database models.
    """
    id = Column(Integer, primary_key=True)

    @declared_attr
    def __tablename__(cls):
        """Autogenerate the tablename for the mapped objects.

        """
        return cls.__name__.lower()

ModelBase = declarative_base(cls=_Base, metaclass=_ModelMeta)

class Dormitory(ModelBase):
    number = Column(String(3), unique=True)
    street = Column(String(20))
    short_name = Column(String(5), unique=True)

    # one to one from Dormitory to VLan
    v_lan = relationship("VLan", uselist=False, backref="dormitory")

class Group(ModelBase):
    name = Column(String(255))

class Host(ModelBase):
    hostname = Column(String(255))

    # one to many from User to Host
    user_id = Column(Integer, ForeignKey("user.id"))

class LogEntry(ModelBase):
    # variably sized string
    message = Column(Text)
    # created
    timestamp = Column(DateTime)

    # many to one from User to LogEntry
    author_id = Column(Integer, ForeignKey("user.id"))

class Membership(ModelBase):
    start_date = Column(DateTime)
    end_date = Column(DateTime)

    user = relationship("User", backref=backref("memberships", order_by=id))
    group = relationship("Group", backref=backref("memberships", order_by=id))

class NetDevice(ModelBase):
    ipv4 = Column(Sring(12), unique=True)
    ipv6 = Column(String(51), unique=True)
    mac = Column(String(13))
    patch_port_id = Column(Integer, ForeignKey("patchport.id"))

class PatchPort(ModelBase):
    name = Column(String(4))

class Rights(ModelBase):
    name = Column(String(255))

class Room(ModelBase):
    number = Column(String(36))
    level = Column(Integer)
    inhabitable = Column(Boolean)
    dormitory_id = Column(Integer, ForeignKey("dormitory.id"))

    dormitory = relationship("Dormitory", backref=backref("rooms",
                                                      order_by=number))

class Subnet(ModelBase):
    address = Column(String(48))
    #address = Column(postgresql.INET)

    # one to one from Subnet to VLan
    v_lan = relationship("VLan", uselist=False, backref="subnet")

class TrafficVolume(ModelBase):
    # 1 (true) for in, 0 (false) for out
    direction = Column(Boolean)
    # how many bytes
    size = Column(BigInteger)
    # when this was logged
    timestamp = Column(DateTime)

    # one to many from User to TrafficVolume
    user_id = Column(Integer, ForeignKey("user.id"))

class User(ModelBase):
    login = Column(String(40))
    name = Column(String(255))
    registration_date = Column(DateTime)

    # one to many from User to Host
    hosts = relationship("Host")
    # one to many from User to LogEntry
    log_entries = relationship("LogEntry")
    #
    room = relationship("Room", backref=backref("users", order_by=id))
    room_id = Column(Integer, ForeignKey("room.id"))
    # one to many from User to TrafficVolume
    traffic_volumes = relationship("TrafficVolume")

class VLan(ModelBase):
    name = Column(String(127))
    tag = Column(Integer)

    # one to one from Dormitory to VLan
    dormitory_id = Column(Integer, ForeignKey("dormitory.id"))
    # one to one from Subnet to VLan
    subnet_id = Column(Integer, ForeignKey("subnet.id"))
<commit_msg>Changed mac in Netdevice to String(12) and added an Comment for MACADDR<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2011 The Pycroft Authors. See the AUTHORS file.
# This file is part of the Pycroft project and licensed under the terms of
# the Apache License, Version 2.0. See the LICENSE file for details.
"""
    pycroft.model.base
    ~~~~~~~~~~~~~~

    This module contains base stuff for all models.

    :copyright: (c) 2011 by AG DSN.
"""
from sqlalchemy import ForeignKey, MetaData
from sqlalchemy import Table, Column
from sqlalchemy.dialects import postgresql
from sqlalchemy.ext.declarative import DeclarativeMeta
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy.orm import relationship, backref
from sqlalchemy.types import BigInteger, Boolean, DateTime, Integer
from sqlalchemy.types import Text, String

_session = None

class _ModelMeta(DeclarativeMeta):
    """Metaclass for all mapped Database objects.
    """
    @property
    def q(cls):
        """This is a shortcut for easy querying of whole objects.

        With this metaclass shortcut you can query a Model with
        Model.q.filter(...) without using the verbose session stuff
        """
        global _session
        if _session is None:
            import pycroft.model.session
            _session = pycroft.model.session.session
        return _session.query(cls)

class _Base(object):
    """Baseclass for all database models.
    """
    id = Column(Integer, primary_key=True)

    @declared_attr
    def __tablename__(cls):
        """Autogenerate the tablename for the mapped objects.

        """
        return cls.__name__.lower()

ModelBase = declarative_base(cls=_Base, metaclass=_ModelMeta)

class Dormitory(ModelBase):
    number = Column(String(3), unique=True)
    street = Column(String(20))
    short_name = Column(String(5), unique=True)

    # one to one from Dormitory to VLan
    v_lan = relationship("VLan", uselist=False, backref="dormitory")

class Group(ModelBase):
    name = Column(String(255))

class Host(ModelBase):
    hostname = Column(String(255))

    # one to many from User to Host
    user_id = Column(Integer, ForeignKey("user.id"))

class LogEntry(ModelBase):
    # variably sized string
    message = Column(Text)
    # created
    timestamp = Column(DateTime)

    # many to one from User to LogEntry
    author_id = Column(Integer, ForeignKey("user.id"))

class Membership(ModelBase):
    start_date = Column(DateTime)
    end_date = Column(DateTime)

    user = relationship("User", backref=backref("memberships", order_by=id))
    group = relationship("Group", backref=backref("memberships", order_by=id))

class NetDevice(ModelBase):
    ipv4 = Column(Sring(12), unique=True)
    ipv6 = Column(String(51), unique=True)
    mac = Column(String(12))
    #mac = Column(postgresql.MACADDR)
    patch_port_id = Column(Integer, ForeignKey("patchport.id"))

class PatchPort(ModelBase):
    name = Column(String(4))

class Rights(ModelBase):
    name = Column(String(255))

class Room(ModelBase):
    number = Column(String(36))
    level = Column(Integer)
    inhabitable = Column(Boolean)
    dormitory_id = Column(Integer, ForeignKey("dormitory.id"))

    dormitory = relationship("Dormitory", backref=backref("rooms",
                                                      order_by=number))

class Subnet(ModelBase):
    address = Column(String(48))
    #address = Column(postgresql.INET)

    # one to one from Subnet to VLan
    v_lan = relationship("VLan", uselist=False, backref="subnet")

class TrafficVolume(ModelBase):
    # 1 (true) for in, 0 (false) for out
    direction = Column(Boolean)
    # how many bytes
    size = Column(BigInteger)
    # when this was logged
    timestamp = Column(DateTime)

    # one to many from User to TrafficVolume
    user_id = Column(Integer, ForeignKey("user.id"))

class User(ModelBase):
    login = Column(String(40))
    name = Column(String(255))
    registration_date = Column(DateTime)

    # one to many from User to Host
    hosts = relationship("Host")
    # one to many from User to LogEntry
    log_entries = relationship("LogEntry")
    #
    room = relationship("Room", backref=backref("users", order_by=id))
    room_id = Column(Integer, ForeignKey("room.id"))
    # one to many from User to TrafficVolume
    traffic_volumes = relationship("TrafficVolume")

class VLan(ModelBase):
    name = Column(String(127))
    tag = Column(Integer)

    # one to one from Dormitory to VLan
    dormitory_id = Column(Integer, ForeignKey("dormitory.id"))
    # one to one from Subnet to VLan
    subnet_id = Column(Integer, ForeignKey("subnet.id"))
<|endoftext|>


<commit_before>"""Contains the emulator itself."""

import time

from .state import State
from .assembly_analysis import Analyzer

class Halt(Exception):
    pass
class InfiniteLoop(Exception):
    pass
class SelfLoop(Exception):
    """Raised when an instruction jumps to itself."""

def format_cell(x):
    """Formats the content of a cell so it is displayed nicely."""
    r = hex(x)[2:]
    if len(r) == 1:
        r = '0' + r
    return r.upper()

class Emulator:
    """Emulator of DigMIPS code."""
    def __init__(self, program, state=None, infinite_loop=False,
            trace_inst=False, trace_mem=False, trace_stack=False):
        self.trace_inst = trace_inst
        self.trace_mem = trace_mem
        self.trace_stack = trace_stack
        self.program = program
        self.state = state or State()
        self.previous_states = set()
        self.detect_same_config = infinite_loop
        self.analyzer = Analyzer(program)
        self.analyzer.analyze()

    def show_trace(self):
        """If the options ask for it, displays the trace."""
        inst = self.program[self.state.pc]
        if self.trace_mem:
            print('Registers: %r' % self.state.registers)
            print('Data: %s' % ' '.join(map(format_cell, self.state.data)))
        if self.trace_inst:
            print('%.03d: %s' % (self.state.pc, inst))
        if self.trace_stack:
            if self.state.pc in self.analyzer.pushes:
                id = self.analyzer.pushes[self.state.pc].arg
                print('Pushed %d (r%d)' % (self.state.registers[id], id))
            if self.state.pc in self.analyzer.pops:
                id = self.analyzer.pops[self.state.pc].arg
                print('Poped %d (r%d)' %
                    (self.analyzer.get_stack_top(self.state), id))

    def run_one(self):
        """Run one instruction."""
        if self.state.pc >= len(self.program):
            raise Halt()
        inst = self.program[self.state.pc]
        self.show_trace()
        old_pc = self.state.pc
        self.state.pc += 1
        inst(self.state)
        self.detect_loops(old_pc)

    def detect_loops(self, old_pc):
        """Determines if we are in a self-loop (using the old_pc) and/or
        in an infinite loop (if the options ask for it)."""
        if self.state.pc == old_pc:
            raise SelfLoop()
        if self.detect_same_config:
            if self.state.freeze() in self.previous_states:
                raise InfiniteLoop()
            self.previous_states.add(self.state.freeze())

    def run(self, max_steps):
        """Run for a certain number of steps."""
        for x in range(0, max_steps):
            self.run_one()

    def run_all(self):
        """Run until we reach the end of the program."""
        try:
            while True:
                self.run_one()
        except Halt:
            pass
<commit_msg>Improve display of push/pop.<commit_after>"""Contains the emulator itself."""

import time

from .state import State
from .assembly_analysis import Analyzer

class Halt(Exception):
    pass
class InfiniteLoop(Exception):
    pass
class SelfLoop(Exception):
    """Raised when an instruction jumps to itself."""

def format_cell(x):
    """Formats the content of a cell so it is displayed nicely."""
    r = hex(x)[2:]
    if len(r) == 1:
        r = '0' + r
    return r.upper()

class Emulator:
    """Emulator of DigMIPS code."""
    def __init__(self, program, state=None, infinite_loop=False,
            trace_inst=False, trace_mem=False, trace_stack=False):
        self.trace_inst = trace_inst
        self.trace_mem = trace_mem
        self.trace_stack = trace_stack
        self.program = program
        self.state = state or State()
        self.previous_states = set()
        self.detect_same_config = infinite_loop
        self.analyzer = Analyzer(program)
        self.analyzer.analyze()

    def show_trace(self):
        """If the options ask for it, displays the trace."""
        inst = self.program[self.state.pc]
        if self.trace_mem:
            print('Registers: %r' % self.state.registers)
            print('Data: %s' % ' '.join(map(format_cell, self.state.data)))
        if self.trace_inst:
            print('%.03d: %s' % (self.state.pc, inst))
        if self.trace_stack:
            if self.state.pc in self.analyzer.pushes:
                id = self.analyzer.pushes[self.state.pc].arg
                print('%.03d: Pushed %d (r%d)' % (self.state.pc,
                    self.state.registers[id], id))
            if self.state.pc in self.analyzer.pops:
                id = self.analyzer.pops[self.state.pc].arg
                print('%.03d Poped %d (r%d)' % (self.state.pc,
                    self.analyzer.get_stack_top(self.state), id))

    def run_one(self):
        """Run one instruction."""
        if self.state.pc >= len(self.program):
            raise Halt()
        inst = self.program[self.state.pc]
        self.show_trace()
        old_pc = self.state.pc
        self.state.pc += 1
        inst(self.state)
        self.detect_loops(old_pc)

    def detect_loops(self, old_pc):
        """Determines if we are in a self-loop (using the old_pc) and/or
        in an infinite loop (if the options ask for it)."""
        if self.state.pc == old_pc:
            raise SelfLoop()
        if self.detect_same_config:
            if self.state.freeze() in self.previous_states:
                raise InfiniteLoop()
            self.previous_states.add(self.state.freeze())

    def run(self, max_steps):
        """Run for a certain number of steps."""
        for x in range(0, max_steps):
            self.run_one()

    def run_all(self):
        """Run until we reach the end of the program."""
        try:
            while True:
                self.run_one()
        except Halt:
            pass
<|endoftext|>


<commit_before>import sys, re, os, modutil, time, types, cPickle, cStringIO, random

from disco import func, util
from disco.comm import download, json
from disco.error import DiscoError, JobError, CommError
from disco.eventmonitor import EventMonitor
from disco.netstring import encode_netstring_fd, encode_netstring_str, decode_netstring_fd
from disco.task import Task

class Params(object):
        def __init__(self, **kwargs):
                self.__dict__.update(kwargs)

        def __getstate__(self):
                return dict((k, util.pack(v))
                        for k, v in self.__dict__.iteritems()
                                if not k.startswith('_'))

        def __setstate__(self, state):
                for k, v in state.iteritems():
                        self.__dict__[k] = util.unpack(v)

class Stats(object):
        def __init__(self, prof_data):
                import marshal
                self.stats = marshal.loads(prof_data)

        def create_stats(self):
                pass

class Disco(object):
        def __init__(self, host):
                self.host = util.disco_host(host)

        def request(self, url, data = None, redir = False, offset = 0):
                try:
                        return download(self.host + url, data = data, redir = redir, offset = offset)
                except CommError, e:
                        raise DiscoError("Got %s, make sure disco master is running at %s" % (e, self.host))

        def get_config(self):
                return json.loads(self.request('/disco/ctrl/load_config_table'))

        def set_config(self, config):
                response = json.loads(self.request('/disco/ctrl/save_config_table', json.dumps(config)))
                if response != 'table saved!':
                        raise DiscoError(response)

        config = property(get_config, set_config)

        def nodeinfo(self):
                return json.loads(self.request('/disco/ctrl/nodeinfo'))

        def joblist(self):
                return json.loads(self.request("/disco/ctrl/joblist"))

        def oob_get(self, name, key):
                try:
                        return util.load_oob(self.host, name, key)
                except CommError, e:
                        if e.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise

        def oob_list(self, name):
                r = self.request("/disco/ctrl/oob_list?name=%s" % name, redir=True)
                return json.loads(r)

        def profile_stats(self, name, mode=''):
                prefix = 'profile-%s' % mode
                f = [s for s in self.oob_list(name) if s.startswith(prefix)]
                if not f:
                        raise JobError("No profile data", self.host, name)

                import pstats
                stats = pstats.Stats(Stats(self.oob_get(name, f[0])))
                for s in f[1:]:
                        stats.add(Stats(self.oob_get(name, s)))
                return stats

        def new_job(self, **kwargs):
                return Job(self, **kwargs)

        def kill(self, name):
                self.request('/disco/ctrl/kill_job', '"%s"' % name)

        def clean(self, name):
                self.request('/disco/ctrl/clean_job', '"%s"' % name)

        def purge(self, name):
                self.request('/disco/ctrl/purge_job', '"%s"' % name)

        def jobspec(self, name):
                r = self.request('/disco/ctrl/parameters?name=%s' % name, redir=True)
                return decode_netstring_fd(cStringIO.StringIO(r))

        def events(self, name, offset = 0):
                def event_iter(events):
                        offs = offset
                        lines = events.splitlines()
                        for i, line in enumerate(lines):
                                if len(line):
                                        offs += len(line) + 1
                                        try:
                                                event = tuple(json.loads(line))
                                        except ValueError:
                                                break
                                        # HTTP range request doesn't like empty ranges:
                                        # Let's ensure that at least the last newline
                                        # is always retrieved.
                                        if i == len(lines) - 1:
                                                offs -= 1
                                        yield offs, event

                r = self.request("/disco/ctrl/rawevents?name=%s" % name,
                                 redir = True,
                                 offset = offset)

                if len(r) < 2:
                        return []
                return event_iter(r)

        def results(self, jobspec, timeout = 2000):
                jobspecifier = JobSpecifier(jobspec)
                data         = json.dumps([timeout, list(jobspecifier.jobnames)])
                results      = json.loads(self.request("/disco/ctrl/get_results", data))

                if type(jobspec) == str:
                        return results[0][1]

                others, active = [], []
                for result in results:
                        if result[1][0] == 'active':
                                active.append(result)
                        else:
                                others.append(result)
                return others, active

        def jobinfo(self, name):
                return json.loads(self.request("/disco/ctrl/jobinfo?name=%s" % name))

        def wait(self, name, show = None, poll_interval = 2, timeout = None, clean = False):
                event_monitor = EventMonitor(show, disco=self, name=name)
                start_time    = time.time()
                while True:
                        status, results = self.results(name, timeout = poll_interval * 1000)
                        event_monitor.refresh()
                        if status == 'ready':
                                if clean:
                                        self.clean(name)
                                return results
                        if status != 'active':
                                raise JobError("Job status %s" % status, self.host, name)
                        if timeout and time.time() - start_time > timeout:
                                raise JobError("Timeout", self.host, name)

class JobSpecifier(list):
        def __init__(self, jobspec):
                super(JobSpecifier, self).__init__([jobspec]
                        if type(jobspec) is str else jobspec)

        @property
        def jobnames(self):
                for job in self:
                        if type(job) is str:
                                yield job
                        elif type(job) is list:
                                yield job[0]
                        else:
                                yield job.name

        def __str__(self):
                return '{%s}' % ', '.join(self.jobnames)

class Job(object):

        funs = ["map", "map_init", "reduce_init", "map_reader", "map_writer",\
                "reduce_reader", "reduce_writer", "reduce", "partition",\
                "combiner", "reduce_input_stream", "reduce_output_stream",\
                "map_input_stream", "map_output_stream"]

        defaults = {"name": None,
                    "map": None,
                    "input": None,
                    "map_init": None,
                    "reduce_init": None,
                    "map_reader": func.map_line_reader,
                    "map_writer": func.netstr_writer,
                    "map_input_stream": None,
                    "map_output_stream": None,
                    "reduce_reader": func.netstr_reader,
                    "reduce_writer": func.netstr_writer,
                    "reduce_input_stream": None,
                    "reduce_output_stream": None,
                    "reduce": None,
                    "partition": func.default_partition,
                    "combiner": None,
                    "nr_maps": None,
                    "scheduler": {},
#XXX: nr_reduces default has changed!
                    "nr_reduces": 1,
                    "sort": False,
                    "params": Params(),
                    "mem_sort_limit": 256 * 1024**2,
                    "ext_params": None,
                    "status_interval": 100000,
                    "required_files": [],
                    "required_modules": [],
                    "profile": False}

        mapreduce_functions = ("map_init",
                               "map_reader",
                               "map",
                               "map_writer",
                               "partition",
                               "combiner",
                               "reduce_init",
                               "reduce_reader",
                               "reduce",
                               "reduce_writer")

        proxy_functions = ("kill",
                           "clean",
                           "purge",
                           "jobspec",
                           "results",
                           "jobinfo",
                           "wait",
                           "oob_get",
                           "oob_list",
                           "profile_stats",
                           "events")

        def __init__(self, master, name='', **kwargs):
                self.master = master
                self.name   = name
                self._run(**kwargs)

        def __getattr__(self, attr):
                if attr in self.proxy_functions:
                        from functools import partial
                        return partial(getattr(self.master, attr), self.name)
                raise AttributeError("%r has no attribute %r" % (self, attr))

        def pack_stack(self, kw, req, stream):
                if stream in kw:
                        req[stream] = encode_netstring_str((f.func_name, util.pack(f))
                                                           for f in util.iterify(kw[stream]))

        def _run(self, **kwargs):
                jobargs = util.DefaultDict(self.defaults.__getitem__, kwargs)

                # -- check parameters --

                # Backwards compatibility
                # (fun_map == map, input_files == input)
                if "fun_map" in kwargs:
                        kwargs["map"] = kwargs["fun_map"]

                if "input_files" in kwargs:
                        kwargs["input"] = kwargs["input_files"]

                if "chunked" in kwargs:
                        raise DeprecationWarning("Argument 'chunked' is deprecated")

                if "nr_maps" in kwargs:
                        sys.stderr.write("Warning: nr_maps is deprecated. "
                                         "Use scheduler = {'max_cores': N} instead.\n")
                        sched = d("scheduler").copy()
                        if "max_cores" not in sched:
                                sched["max_cores"] = int(kw["nr_maps"])
                        kw["scheduler"] = sched

                if not "input" in kwargs:
                        raise DiscoError("Argument input is required")

                if not ("map" in kwargs or "reduce" in kwargs):
                        raise DiscoError("Specify map and/or reduce")

                for p in kwargs:
                        if p not in Job.defaults:
                                raise DiscoError("Unknown argument: %s" % p)

                input = kwargs["input"]

                # -- initialize request --

                request = {"prefix": self.name,
                           "version": ".".join(map(str, sys.version_info[:2])),
                           "params": cPickle.dumps(jobargs['params'], cPickle.HIGHEST_PROTOCOL),
                           "sort": str(int(jobargs['sort'])),
                           "mem_sort_limit": str(jobargs['mem_sort_limit']),
                           "status_interval": str(jobargs['status_interval']),
                           "profile": str(int(jobargs['profile']))}

                # -- required modules --

                if 'required_modules' in kwargs:
                        rm = kwargs['required_modules']
                else:
                        functions = util.flatten(util.iterify(jobargs[f])
                                                 for f in self.mapreduce_functions)
                        rm = modutil.find_modules([f for f in functions
                                if isinstance(f, types.FunctionType)])

                send_mod = []
                imp_mod = []
                for mod in rm:
                        if type(mod) == tuple:
                                send_mod.append(mod[1])
                                mod = mod[0]
                        imp_mod.append(mod)

                request["required_modules"] = " ".join(imp_mod)
                rf = util.pack_files(send_mod)

                # -- input & output streams --

                for stream in ["map_input_stream", "map_output_stream",
                               "reduce_input_stream", "reduce_output_stream"]:
                        self.pack_stack(kwargs, request, stream)

                # -- required files --

                if "required_files" in kwargs:
                        if isinstance(kwargs["required_files"], dict):
                                rf.update(kwargs["required_files"])
                        else:
                                rf.update(util.pack_files(\
                                        kwargs["required_files"]))
                if rf:
                        request["required_files"] = util.pack(rf)

                # -- scheduler --

                sched = jobargs["scheduler"]
                sched_keys = ["max_cores", "force_local", "force_remote"]

                if "max_cores" not in sched:
                        sched["max_cores"] = 2**31
                elif sched["max_cores"] < 1:
                        raise DiscoError("max_cores must be >= 1")

                for k in sched_keys:
                        if k in sched:
                                request["sched_" + k] = str(sched[k])

                # -- map --

                if 'map' in kwargs:
                        k = 'ext_map' if isinstance(kwargs['map'], dict) else 'map'
                        request[k] = util.pack(kwargs['map'])

                        for function_name in ('map_init',
                                              'map_reader',
                                              'map_writer',
                                              'partition',
                                              'combiner'):
                                function = jobargs[function_name]
                                if function:
                                        request[function_name] = util.pack(function)

                        def inputlist(input):
                                if hasattr(input, '__iter__'):
                                        return ['\n'.join(reversed(list(input)))]
                                return util.urllist(input)
                        input = [e for i in input for e in inputlist(i)]

                # -- only reduce --

                else:
                        # XXX: Check for redundant inputs, external &
                        # partitioned inputs
                        input = [url for i in input for url in util.urllist(i)]

                request['input'] = ' '.join(input)

                if 'ext_params' in kwargs:
                        e = kwargs['ext_params']
                        request['ext_params'] = encode_netstring_fd(e)
                                if isinstance(e, dict) else e

                # -- reduce --

                nr_reduces = jobargs['nr_reduces']
                if 'reduce' in kwargs:
                        k = 'ext_reduce' if isinstance(kwargs['reduce'], dict) else 'reduce'
                        request[k] = util.pack(kwargs['reduce'])

                        for function_name in ('reduce_reader', 'reduce_writer', 'reduce_init'):
                                function = jobargs[function_name]
                                if function:
                                        request[function_name] = util.pack(function)

                request['nr_reduces'] = str(nr_reduces)

                # -- encode and send the request --

                reply = self.master.request("/disco/job/new", encode_netstring_fd(request))

                if not reply.startswith('job started:'):
                        raise DiscoError("Failed to start a job. Server replied: " + reply)
                self.name = reply.split(':', 1)[1]

def result_iterator(results,
                    notifier = None,
                    reader = func.netstr_reader,
                    input_stream = [func.map_input_stream],
                    params = None):

        task = Task(result_iterator = True)
        for fun in input_stream:
                fun.func_globals.setdefault("Task", task)

        res = []
        res = [url for r in results for url in util.urllist(r)]

        for url in res:
                fd = sze = None
                for fun in input_stream:
                        fd, sze, url = fun(fd, sze, url, params)

                if notifier:
                        notifier(url)

                for x in reader(fd, sze, url):
                        yield x

<commit_msg>Fixed an unexpected indent<commit_after>import sys, re, os, modutil, time, types, cPickle, cStringIO, random

from disco import func, util
from disco.comm import download, json
from disco.error import DiscoError, JobError, CommError
from disco.eventmonitor import EventMonitor
from disco.netstring import encode_netstring_fd, encode_netstring_str, decode_netstring_fd
from disco.task import Task

class Params(object):
        def __init__(self, **kwargs):
                self.__dict__.update(kwargs)

        def __getstate__(self):
                return dict((k, util.pack(v))
                        for k, v in self.__dict__.iteritems()
                                if not k.startswith('_'))

        def __setstate__(self, state):
                for k, v in state.iteritems():
                        self.__dict__[k] = util.unpack(v)

class Stats(object):
        def __init__(self, prof_data):
                import marshal
                self.stats = marshal.loads(prof_data)

        def create_stats(self):
                pass

class Disco(object):
        def __init__(self, host):
                self.host = util.disco_host(host)

        def request(self, url, data = None, redir = False, offset = 0):
                try:
                        return download(self.host + url, data = data, redir = redir, offset = offset)
                except CommError, e:
                        raise DiscoError("Got %s, make sure disco master is running at %s" % (e, self.host))

        def get_config(self):
                return json.loads(self.request('/disco/ctrl/load_config_table'))

        def set_config(self, config):
                response = json.loads(self.request('/disco/ctrl/save_config_table', json.dumps(config)))
                if response != 'table saved!':
                        raise DiscoError(response)

        config = property(get_config, set_config)

        def nodeinfo(self):
                return json.loads(self.request('/disco/ctrl/nodeinfo'))

        def joblist(self):
                return json.loads(self.request("/disco/ctrl/joblist"))

        def oob_get(self, name, key):
                try:
                        return util.load_oob(self.host, name, key)
                except CommError, e:
                        if e.http_code == 404:
                                raise DiscoError("Unknown key or job name")
                        raise

        def oob_list(self, name):
                r = self.request("/disco/ctrl/oob_list?name=%s" % name, redir=True)
                return json.loads(r)

        def profile_stats(self, name, mode=''):
                prefix = 'profile-%s' % mode
                f = [s for s in self.oob_list(name) if s.startswith(prefix)]
                if not f:
                        raise JobError("No profile data", self.host, name)

                import pstats
                stats = pstats.Stats(Stats(self.oob_get(name, f[0])))
                for s in f[1:]:
                        stats.add(Stats(self.oob_get(name, s)))
                return stats

        def new_job(self, **kwargs):
                return Job(self, **kwargs)

        def kill(self, name):
                self.request('/disco/ctrl/kill_job', '"%s"' % name)

        def clean(self, name):
                self.request('/disco/ctrl/clean_job', '"%s"' % name)

        def purge(self, name):
                self.request('/disco/ctrl/purge_job', '"%s"' % name)

        def jobspec(self, name):
                r = self.request('/disco/ctrl/parameters?name=%s' % name, redir=True)
                return decode_netstring_fd(cStringIO.StringIO(r))

        def events(self, name, offset = 0):
                def event_iter(events):
                        offs = offset
                        lines = events.splitlines()
                        for i, line in enumerate(lines):
                                if len(line):
                                        offs += len(line) + 1
                                        try:
                                                event = tuple(json.loads(line))
                                        except ValueError:
                                                break
                                        # HTTP range request doesn't like empty ranges:
                                        # Let's ensure that at least the last newline
                                        # is always retrieved.
                                        if i == len(lines) - 1:
                                                offs -= 1
                                        yield offs, event

                r = self.request("/disco/ctrl/rawevents?name=%s" % name,
                                 redir = True,
                                 offset = offset)

                if len(r) < 2:
                        return []
                return event_iter(r)

        def results(self, jobspec, timeout = 2000):
                jobspecifier = JobSpecifier(jobspec)
                data         = json.dumps([timeout, list(jobspecifier.jobnames)])
                results      = json.loads(self.request("/disco/ctrl/get_results", data))

                if type(jobspec) == str:
                        return results[0][1]

                others, active = [], []
                for result in results:
                        if result[1][0] == 'active':
                                active.append(result)
                        else:
                                others.append(result)
                return others, active

        def jobinfo(self, name):
                return json.loads(self.request("/disco/ctrl/jobinfo?name=%s" % name))

        def wait(self, name, show = None, poll_interval = 2, timeout = None, clean = False):
                event_monitor = EventMonitor(show, disco=self, name=name)
                start_time    = time.time()
                while True:
                        status, results = self.results(name, timeout = poll_interval * 1000)
                        event_monitor.refresh()
                        if status == 'ready':
                                if clean:
                                        self.clean(name)
                                return results
                        if status != 'active':
                                raise JobError("Job status %s" % status, self.host, name)
                        if timeout and time.time() - start_time > timeout:
                                raise JobError("Timeout", self.host, name)

class JobSpecifier(list):
        def __init__(self, jobspec):
                super(JobSpecifier, self).__init__([jobspec]
                        if type(jobspec) is str else jobspec)

        @property
        def jobnames(self):
                for job in self:
                        if type(job) is str:
                                yield job
                        elif type(job) is list:
                                yield job[0]
                        else:
                                yield job.name

        def __str__(self):
                return '{%s}' % ', '.join(self.jobnames)

class Job(object):

        funs = ["map", "map_init", "reduce_init", "map_reader", "map_writer",\
                "reduce_reader", "reduce_writer", "reduce", "partition",\
                "combiner", "reduce_input_stream", "reduce_output_stream",\
                "map_input_stream", "map_output_stream"]

        defaults = {"name": None,
                    "map": None,
                    "input": None,
                    "map_init": None,
                    "reduce_init": None,
                    "map_reader": func.map_line_reader,
                    "map_writer": func.netstr_writer,
                    "map_input_stream": None,
                    "map_output_stream": None,
                    "reduce_reader": func.netstr_reader,
                    "reduce_writer": func.netstr_writer,
                    "reduce_input_stream": None,
                    "reduce_output_stream": None,
                    "reduce": None,
                    "partition": func.default_partition,
                    "combiner": None,
                    "nr_maps": None,
                    "scheduler": {},
#XXX: nr_reduces default has changed!
                    "nr_reduces": 1,
                    "sort": False,
                    "params": Params(),
                    "mem_sort_limit": 256 * 1024**2,
                    "ext_params": None,
                    "status_interval": 100000,
                    "required_files": [],
                    "required_modules": [],
                    "profile": False}

        mapreduce_functions = ("map_init",
                               "map_reader",
                               "map",
                               "map_writer",
                               "partition",
                               "combiner",
                               "reduce_init",
                               "reduce_reader",
                               "reduce",
                               "reduce_writer")

        proxy_functions = ("kill",
                           "clean",
                           "purge",
                           "jobspec",
                           "results",
                           "jobinfo",
                           "wait",
                           "oob_get",
                           "oob_list",
                           "profile_stats",
                           "events")

        def __init__(self, master, name='', **kwargs):
                self.master = master
                self.name   = name
                self._run(**kwargs)

        def __getattr__(self, attr):
                if attr in self.proxy_functions:
                        from functools import partial
                        return partial(getattr(self.master, attr), self.name)
                raise AttributeError("%r has no attribute %r" % (self, attr))

        def pack_stack(self, kw, req, stream):
                if stream in kw:
                        req[stream] = encode_netstring_str((f.func_name, util.pack(f))
                                                           for f in util.iterify(kw[stream]))

        def _run(self, **kwargs):
                jobargs = util.DefaultDict(self.defaults.__getitem__, kwargs)

                # -- check parameters --

                # Backwards compatibility
                # (fun_map == map, input_files == input)
                if "fun_map" in kwargs:
                        kwargs["map"] = kwargs["fun_map"]

                if "input_files" in kwargs:
                        kwargs["input"] = kwargs["input_files"]

                if "chunked" in kwargs:
                        raise DeprecationWarning("Argument 'chunked' is deprecated")

                if "nr_maps" in kwargs:
                        sys.stderr.write("Warning: nr_maps is deprecated. "
                                         "Use scheduler = {'max_cores': N} instead.\n")
                        sched = d("scheduler").copy()
                        if "max_cores" not in sched:
                                sched["max_cores"] = int(kw["nr_maps"])
                        kw["scheduler"] = sched

                if not "input" in kwargs:
                        raise DiscoError("Argument input is required")

                if not ("map" in kwargs or "reduce" in kwargs):
                        raise DiscoError("Specify map and/or reduce")

                for p in kwargs:
                        if p not in Job.defaults:
                                raise DiscoError("Unknown argument: %s" % p)

                input = kwargs["input"]

                # -- initialize request --

                request = {"prefix": self.name,
                           "version": ".".join(map(str, sys.version_info[:2])),
                           "params": cPickle.dumps(jobargs['params'], cPickle.HIGHEST_PROTOCOL),
                           "sort": str(int(jobargs['sort'])),
                           "mem_sort_limit": str(jobargs['mem_sort_limit']),
                           "status_interval": str(jobargs['status_interval']),
                           "profile": str(int(jobargs['profile']))}

                # -- required modules --

                if 'required_modules' in kwargs:
                        rm = kwargs['required_modules']
                else:
                        functions = util.flatten(util.iterify(jobargs[f])
                                                 for f in self.mapreduce_functions)
                        rm = modutil.find_modules([f for f in functions
                                if isinstance(f, types.FunctionType)])

                send_mod = []
                imp_mod = []
                for mod in rm:
                        if type(mod) == tuple:
                                send_mod.append(mod[1])
                                mod = mod[0]
                        imp_mod.append(mod)

                request["required_modules"] = " ".join(imp_mod)
                rf = util.pack_files(send_mod)

                # -- input & output streams --

                for stream in ["map_input_stream", "map_output_stream",
                               "reduce_input_stream", "reduce_output_stream"]:
                        self.pack_stack(kwargs, request, stream)

                # -- required files --

                if "required_files" in kwargs:
                        if isinstance(kwargs["required_files"], dict):
                                rf.update(kwargs["required_files"])
                        else:
                                rf.update(util.pack_files(\
                                        kwargs["required_files"]))
                if rf:
                        request["required_files"] = util.pack(rf)

                # -- scheduler --

                sched = jobargs["scheduler"]
                sched_keys = ["max_cores", "force_local", "force_remote"]

                if "max_cores" not in sched:
                        sched["max_cores"] = 2**31
                elif sched["max_cores"] < 1:
                        raise DiscoError("max_cores must be >= 1")

                for k in sched_keys:
                        if k in sched:
                                request["sched_" + k] = str(sched[k])

                # -- map --

                if 'map' in kwargs:
                        k = 'ext_map' if isinstance(kwargs['map'], dict) else 'map'
                        request[k] = util.pack(kwargs['map'])

                        for function_name in ('map_init',
                                              'map_reader',
                                              'map_writer',
                                              'partition',
                                              'combiner'):
                                function = jobargs[function_name]
                                if function:
                                        request[function_name] = util.pack(function)

                        def inputlist(input):
                                if hasattr(input, '__iter__'):
                                        return ['\n'.join(reversed(list(input)))]
                                return util.urllist(input)
                        input = [e for i in input for e in inputlist(i)]

                # -- only reduce --

                else:
                        # XXX: Check for redundant inputs, external &
                        # partitioned inputs
                        input = [url for i in input for url in util.urllist(i)]

                request['input'] = ' '.join(input)

                if 'ext_params' in kwargs:
                        e = kwargs['ext_params']
                        request['ext_params'] = encode_netstring_fd(e)\
                                if isinstance(e, dict) else e

                # -- reduce --

                nr_reduces = jobargs['nr_reduces']
                if 'reduce' in kwargs:
                        k = 'ext_reduce' if isinstance(kwargs['reduce'], dict) else 'reduce'
                        request[k] = util.pack(kwargs['reduce'])

                        for function_name in ('reduce_reader', 'reduce_writer', 'reduce_init'):
                                function = jobargs[function_name]
                                if function:
                                        request[function_name] = util.pack(function)

                request['nr_reduces'] = str(nr_reduces)

                # -- encode and send the request --

                reply = self.master.request("/disco/job/new", encode_netstring_fd(request))

                if not reply.startswith('job started:'):
                        raise DiscoError("Failed to start a job. Server replied: " + reply)
                self.name = reply.split(':', 1)[1]

def result_iterator(results,
                    notifier = None,
                    reader = func.netstr_reader,
                    input_stream = [func.map_input_stream],
                    params = None):

        task = Task(result_iterator = True)
        for fun in input_stream:
                fun.func_globals.setdefault("Task", task)

        res = []
        res = [url for r in results for url in util.urllist(r)]

        for url in res:
                fd = sze = None
                for fun in input_stream:
                        fd, sze, url = fun(fd, sze, url, params)

                if notifier:
                        notifier(url)

                for x in reader(fd, sze, url):
                        yield x

<|endoftext|>