<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals
import os
import sys

on_rtd = os.environ.get('READTHEDOCS', None) == 'True'

# http://read-the-docs.readthedocs.org/en/latest/faq.html#i-get-import-errors-on-libraries-that-depend-on-c-modules
MOCK_MODULES = ['pytz']
if MOCK_MODULES and on_rtd:
    if sys.version_info > (3, 3):
        from unittest.mock import MagicMock
    else:
        from mock import Mock as MagicMock

    class Mock(MagicMock):
        @classmethod
        def __getattr__(cls, name):
            return Mock()

    sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)

# -- General configuration ----------------------------------------------------

project = 'Astral'
author = 'Simon Kennedy'
copyright = '2009-2018, %s' % author
version = '1.5'
release = '1.5'

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
extensions = ['sphinx.ext.autodoc', 'sphinx.ext.intersphinx']
intersphinx_mapping = {'python': ('http://docs.python.org/3', None)}

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
sys.path.insert(0, os.path.join(os.path.abspath('..')))

# The suffix of source filenames.
source_suffix = '.rst'

# The encoding of source files.
#source_encoding = 'utf-8'

# The master toctree document.
master_doc = 'index'

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#today = ''
# Else, today_fmt is used as the format for a strftime call.
#today_fmt = '%B %d, %Y'

# List of documents that shouldn't be included in the build.
#unused_docs = []

# List of directories, relative to source directory, that shouldn't be searched
# for source files.
exclude_trees = []

# The reST default role (used for this markup: `text`) to use for all documents.
#default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
#add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'

# A list of ignored prefixes for module index sorting.
#modindex_common_prefix = []

# Add any paths that contain templates here, relative to this directory.
templates_path = ['templates']

# -- Options for HTML output ---------------------------------------------------

if not on_rtd:
    project_home = os.environ.get('PROJECT_HOME', None)
    if not project_home:
        dev_home = os.environ.get('DEV_HOME', None)
        if dev_home:
            project_home = os.path.join(os.path.expanduser(dev_home), 'projects')
    else:
        project_home = os.path.expanduser(project_home)

    if project_home:
        theme_root = os.path.join(project_home, 'sphinx-theme', 'sffjunkie', 'trunk')
        html_theme_path = [theme_root]
    else:
        raise OSError('Unable to find theme root: Please set the PROJECT_HOME environment variable')

    html_theme = 'sffjunkie'
    html_theme_options = {'logo_shadow': True, 'fixed_header': False}
else:
    html_theme = 'default'

# The name for this set of Sphinx documents.  If None, it defaults to
# "<project> v<release> documentation".
#html_title = None

# A shorter title for the navigation bar.  Default is the same as html_title.
#html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
html_logo = os.path.join('static', 'earth_sun.png')

# The name of an image file (within the static path) to use as favicon of the
# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
#if not on_rtd:
#    html_favicon = os.path.join('_static', 'favicon.png')

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['static']

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
#html_last_updated_fmt = '%b %d, %Y'

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
#html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
#html_sidebars = {}

# Additional templates that should be rendered to pages, maps page names to
# template names.
#html_additional_pages = {}

# If false, no module index is generated.
html_use_modindex = False

# If false, no index is generated.
#html_use_index = True

# If true, the index is split into individual pages for each letter.
#html_split_index = False

# If true, links to the reST sources are added to the pages.
#html_show_sourcelink = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
#html_use_opensearch = ''

# If nonempty, this is the file name suffix for HTML files (e.g. ".xhtml").
#html_file_suffix = ''

# -- Options for HTML Help output --------------------------------------------------

# Output file base name for HTML help builder.
htmlhelp_basename = 'AstralDoc'

# -- Options for LaTeX output --------------------------------------------------

# The paper size ('letter' or 'a4').
#latex_paper_size = 'a4'
latex_elements = {'papersize': 'a4'}

# The font size ('10pt', '11pt' or '12pt').
#latex_font_size = '10pt'

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, documentclass [howto/manual]).
latex_documents = [
  ('index', 'Astral.tex', 'Astral v%s' % release, author, 'manual'),
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
#latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
#latex_use_parts = False

# Additional stuff for the LaTeX preamble.
#latex_preamble = ''

# Documents to append as an appendix to all manuals.
#latex_appendices = []

# If false, no module index is generated.
#latex_use_modindex = True
<commit_msg>html_use_modindex is deprecated<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals
import os
import sys

on_rtd = os.environ.get('READTHEDOCS', None) == 'True'

# http://read-the-docs.readthedocs.org/en/latest/faq.html#i-get-import-errors-on-libraries-that-depend-on-c-modules
MOCK_MODULES = ['pytz']
if MOCK_MODULES and on_rtd:
    if sys.version_info > (3, 3):
        from unittest.mock import MagicMock
    else:
        from mock import Mock as MagicMock

    class Mock(MagicMock):
        @classmethod
        def __getattr__(cls, name):
            return Mock()

    sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)

# -- General configuration ----------------------------------------------------

project = 'Astral'
author = 'Simon Kennedy'
copyright = '2009-2018, %s' % author
version = '1.5'
release = '1.5'

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
extensions = ['sphinx.ext.autodoc', 'sphinx.ext.intersphinx']
intersphinx_mapping = {'python': ('http://docs.python.org/3', None)}

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
sys.path.insert(0, os.path.join(os.path.abspath('..')))

# The suffix of source filenames.
source_suffix = '.rst'

# The encoding of source files.
#source_encoding = 'utf-8'

# The master toctree document.
master_doc = 'index'

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#today = ''
# Else, today_fmt is used as the format for a strftime call.
#today_fmt = '%B %d, %Y'

# List of documents that shouldn't be included in the build.
#unused_docs = []

# List of directories, relative to source directory, that shouldn't be searched
# for source files.
exclude_trees = []

# The reST default role (used for this markup: `text`) to use for all documents.
#default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
#add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'

# A list of ignored prefixes for module index sorting.
#modindex_common_prefix = []

# Add any paths that contain templates here, relative to this directory.
templates_path = ['templates']

# -- Options for HTML output ---------------------------------------------------

if not on_rtd:
    project_home = os.environ.get('PROJECT_HOME', None)
    if not project_home:
        dev_home = os.environ.get('DEV_HOME', None)
        if dev_home:
            project_home = os.path.join(os.path.expanduser(dev_home), 'projects')
    else:
        project_home = os.path.expanduser(project_home)

    if project_home:
        theme_root = os.path.join(project_home, 'sphinx-theme', 'sffjunkie', 'trunk')
        html_theme_path = [theme_root]
    else:
        raise OSError('Unable to find theme root: Please set the PROJECT_HOME environment variable')

    html_theme = 'sffjunkie'
    html_theme_options = {'logo_shadow': True, 'fixed_header': False}
else:
    html_theme = 'default'

# The name for this set of Sphinx documents.  If None, it defaults to
# "<project> v<release> documentation".
#html_title = None

# A shorter title for the navigation bar.  Default is the same as html_title.
#html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
html_logo = os.path.join('static', 'earth_sun.png')

# The name of an image file (within the static path) to use as favicon of the
# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
#if not on_rtd:
#    html_favicon = os.path.join('_static', 'favicon.png')

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['static']

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
#html_last_updated_fmt = '%b %d, %Y'

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
#html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
#html_sidebars = {}

# Additional templates that should be rendered to pages, maps page names to
# template names.
#html_additional_pages = {}

# If false, no module index is generated.
html_domain_indices = False

# If false, no index is generated.
#html_use_index = True

# If true, the index is split into individual pages for each letter.
#html_split_index = False

# If true, links to the reST sources are added to the pages.
#html_show_sourcelink = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
#html_use_opensearch = ''

# If nonempty, this is the file name suffix for HTML files (e.g. ".xhtml").
#html_file_suffix = ''

# -- Options for HTML Help output --------------------------------------------------

# Output file base name for HTML help builder.
htmlhelp_basename = 'AstralDoc'

# -- Options for LaTeX output --------------------------------------------------

# The paper size ('letter' or 'a4').
#latex_paper_size = 'a4'
latex_elements = {'papersize': 'a4'}

# The font size ('10pt', '11pt' or '12pt').
#latex_font_size = '10pt'

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, documentclass [howto/manual]).
latex_documents = [
  ('index', 'Astral.tex', 'Astral v%s' % release, author, 'manual'),
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
#latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
#latex_use_parts = False

# Additional stuff for the LaTeX preamble.
#latex_preamble = ''

# Documents to append as an appendix to all manuals.
#latex_appendices = []

# If false, no module index is generated.
#latex_use_modindex = True
<|endoftext|>


<commit_before>"""
This module defines exporters for export of schedules and Graphs to SVG.
"""

import manuallabour.exporters.common as common
import pygraphviz as pgv
import os
import tempfile

def graphviz_add_obj_edges(graph,s_id,objs,**kwargs):
    """
    add objects in the list objs to graph.

    kwargs:
    attr: dict of edge attributes
    opt: additional edge attributes for optional objects
    res: also add links to resources referenced by this object
    """
    for obj in objs.values():
        o_id = 'o_' + obj["obj_id"]
        if obj["optional"] and "opt" in kwargs:
            kwargs["attr"].update(kwargs["opt"])
        kwargs["attr"]["label"] = obj["quantity"]
        if obj["created"]:
            graph.add_edge(s_id,o_id,**kwargs["attr"])
        else:
            graph.add_edge(o_id,s_id,**kwargs["attr"])
        if kwargs["res"]:
            for img in obj["images"]:
                graph.add_edge('r_' + img["blob_id"],o_id)

class GraphSVGExporter(common.GraphExporterBase):
    """
    Exporter to export graphs to svg files.
    """
    def __init__(self,with_objects=False,with_resources=False):
        common.GraphExporterBase.__init__(self)
        self.with_objects = with_objects
        self.with_resources = with_resources

    def export(self,graph,store,path,**kwargs):
        common.GraphExporterBase.export(self,graph,store,path,**kwargs)

        result = pgv.AGraph(directed=True)

        steps = {}
        for ref in graph.steps:
            steps[ref.step_id] = ref.dereference(store)

        #Nodes
        for alias,step_dict in steps.iteritems():
            s_id = "s_" + alias
            result.add_node(s_id,label=step_dict["title"])

        if self.with_objects:
            for o_id, obj in store.iter_obj():
                o_id = 'o_' + o_id
                result.add_node(o_id,label=obj.name,shape='rectangle')

        if self.with_resources:
            for blob_id in store.iter_blob():
                r_id = 'r_' + blob_id
                result.add_node(r_id,label=blob_id[:6],shape='diamond')

        #Edges
        for alias,children in graph.children.iteritems():
            for child in children:
                result.add_edge('s_' + alias,'s_' + child)

        if self.with_objects:
            for alias,step_dict in steps.iteritems():
                s_id = 's_' + alias

                args = dict(
                    attr={'color' : 'blue'},
                    opt={'style' : 'dashed'},
                    res=self.with_resources
                )
                graphviz_add_obj_edges(result,s_id,step_dict["parts"],**args)

                args["attr"] = {'color' : 'red'}
                graphviz_add_obj_edges(result,s_id,step_dict["tools"],**args)

                args["attr"] = {'color' : 'brown'}
                graphviz_add_obj_edges(result,s_id,step_dict["results"],**args)

        if self.with_resources:
            for alias,step_dict in steps.iteritems():
                s_id = 's_' + alias

                for res in step_dict["files"].values():
                    result.add_edge('r_' + res["blob_id"],s_id,color='orange')
                for res in step_dict["images"].values():
                    result.add_edge('r_' + res["blob_id"],s_id,color='green')

        result.draw(path,prog='dot')

    def render(self,graph,store,**kwargs):
        common.GraphExporterBase.render(self,graph,store,**kwargs)

        _filehandle,filename = tempfile.mkstemp(suffix=".svg")
        self.export(graph,store,filename,**kwargs)
        result = open(filename).read()
        os.remove(filename)
        return result

class ScheduleSVGExporter(common.ScheduleExporterBase):
    """
    Exporter to export schedules to svg files.
    """
    def __init__(self,with_objects=False,with_resources=False):
        common.ScheduleExporterBase.__init__(self)
        self.with_objects = with_objects
        self.with_resources = with_resources

    def export(self,schedule,store,path,**kwargs):
        common.ScheduleExporterBase.export(self,schedule,store,path,**kwargs)

        result = pgv.AGraph(directed=True)

        #Nodes
        for ref in schedule.steps:
            s_id = 's_' + str(ref.step_nr)
            step_dict = ref.dereference(store)
            result.add_node(s_id,label=step_dict["title"])
            if ref.step_nr > 1:
                result.add_edge('s_' + str(ref.step_nr - 1),s_id)

        if self.with_objects:
            for o_id, obj in store.iter_obj():
                o_id = 'o_' + o_id
                result.add_node(o_id,label=obj.name,shape='rectangle')

        if self.with_resources:
            for blob_id in store.iter_blob():
                result.add_node(blob_id,label=blob_id[:6],shape='diamond')

        #edges
        for ref in schedule.steps:
            s_id = 's_' + str(ref.step_nr)
            if ref.step_nr > 1:
                result.add_edge('s_' + str(ref.step_nr - 1),s_id)

        if self.with_objects:
            for ref in schedule.steps:
                s_id = 's_' + str(ref.step_nr)
                step_dict = ref.dereference(store)

                args = dict(
                    attr={'color' : 'blue'},
                    opt={'style' : 'dashed'},
                    res=self.with_resources
                )
                graphviz_add_obj_edges(result,s_id,step_dict["parts"],**args)

                args["attr"] = {'color' : 'red'}
                graphviz_add_obj_edges(result,s_id,step_dict["tools"],**args)

                args["attr"] = {'color' : 'brown'}
                graphviz_add_obj_edges(result,s_id,step_dict["results"],**args)

        if self.with_resources:
            for ref in schedule.steps:
                s_id = 's_' + str(ref.step_nr)
                step = store.get_step(ref.step_id)

                for fil in step.files.values():
                    result.add_edge('r_' + fil.blob_id,s_id,color='orange')
                for img in step.images.values():
                    result.add_edge('r_' + img.blob_id,s_id,color='green')

        result.draw(path,prog='dot')

<commit_msg>non-strict graphs for exporters<commit_after>"""
This module defines exporters for export of schedules and Graphs to SVG.
"""

import manuallabour.exporters.common as common
import pygraphviz as pgv
import os
import tempfile

def graphviz_add_obj_edges(graph,s_id,objs,**kwargs):
    """
    add objects in the list objs to graph.

    kwargs:
    attr: dict of edge attributes
    opt: additional edge attributes for optional objects
    res: also add links to resources referenced by this object
    """
    for obj in objs.values():
        o_id = 'o_' + obj["obj_id"]
        if obj["optional"] and "opt" in kwargs:
            kwargs["attr"].update(kwargs["opt"])
        kwargs["attr"]["label"] = obj["quantity"]
        if obj["created"]:
            graph.add_edge(s_id,o_id,**kwargs["attr"])
        else:
            graph.add_edge(o_id,s_id,**kwargs["attr"])
        if kwargs["res"]:
            for img in obj["images"]:
                graph.add_edge('r_' + img["blob_id"],o_id)

class GraphSVGExporter(common.GraphExporterBase):
    """
    Exporter to export graphs to svg files.
    """
    def __init__(self,with_objects=False,with_resources=False):
        common.GraphExporterBase.__init__(self)
        self.with_objects = with_objects
        self.with_resources = with_resources

    def export(self,graph,store,path,**kwargs):
        common.GraphExporterBase.export(self,graph,store,path,**kwargs)

        result = pgv.AGraph(directed=True,strict=False)

        steps = {}
        for ref in graph.steps:
            steps[ref.step_id] = ref.dereference(store)

        #Nodes
        for alias,step_dict in steps.iteritems():
            s_id = "s_" + alias
            result.add_node(s_id,label=step_dict["title"])

        if self.with_objects:
            for o_id, obj in store.iter_obj():
                o_id = 'o_' + o_id
                result.add_node(o_id,label=obj.name,shape='rectangle')

        if self.with_resources:
            for blob_id in store.iter_blob():
                r_id = 'r_' + blob_id
                result.add_node(r_id,label=blob_id[:6],shape='diamond')

        #Edges
        for alias,children in graph.children.iteritems():
            for child in children:
                result.add_edge('s_' + alias,'s_' + child)

        if self.with_objects:
            for alias,step_dict in steps.iteritems():
                s_id = 's_' + alias

                args = dict(
                    attr={'color' : 'blue'},
                    opt={'style' : 'dashed'},
                    res=self.with_resources
                )
                graphviz_add_obj_edges(result,s_id,step_dict["parts"],**args)

                args["attr"] = {'color' : 'red'}
                graphviz_add_obj_edges(result,s_id,step_dict["tools"],**args)

                args["attr"] = {'color' : 'brown'}
                graphviz_add_obj_edges(result,s_id,step_dict["results"],**args)

        if self.with_resources:
            for alias,step_dict in steps.iteritems():
                s_id = 's_' + alias

                for res in step_dict["files"].values():
                    result.add_edge('r_' + res["blob_id"],s_id,color='orange')
                for res in step_dict["images"].values():
                    result.add_edge('r_' + res["blob_id"],s_id,color='green')

        result.draw(path,prog='dot')

    def render(self,graph,store,**kwargs):
        common.GraphExporterBase.render(self,graph,store,**kwargs)

        _filehandle,filename = tempfile.mkstemp(suffix=".svg")
        self.export(graph,store,filename,**kwargs)
        result = open(filename).read()
        os.remove(filename)
        return result

class ScheduleSVGExporter(common.ScheduleExporterBase):
    """
    Exporter to export schedules to svg files.
    """
    def __init__(self,with_objects=False,with_resources=False):
        common.ScheduleExporterBase.__init__(self)
        self.with_objects = with_objects
        self.with_resources = with_resources

    def export(self,schedule,store,path,**kwargs):
        common.ScheduleExporterBase.export(self,schedule,store,path,**kwargs)

        result = pgv.AGraph(directed=True,strict=False)

        #Nodes
        for ref in schedule.steps:
            s_id = 's_' + str(ref.step_nr)
            step_dict = ref.dereference(store)
            result.add_node(s_id,label=step_dict["title"])
            if ref.step_nr > 1:
                result.add_edge('s_' + str(ref.step_nr - 1),s_id)

        if self.with_objects:
            for o_id, obj in store.iter_obj():
                o_id = 'o_' + o_id
                result.add_node(o_id,label=obj.name,shape='rectangle')

        if self.with_resources:
            for blob_id in store.iter_blob():
                result.add_node(blob_id,label=blob_id[:6],shape='diamond')

        #edges
        for ref in schedule.steps:
            s_id = 's_' + str(ref.step_nr)
            if ref.step_nr > 1:
                result.add_edge('s_' + str(ref.step_nr - 1),s_id)

        if self.with_objects:
            for ref in schedule.steps:
                s_id = 's_' + str(ref.step_nr)
                step_dict = ref.dereference(store)

                args = dict(
                    attr={'color' : 'blue'},
                    opt={'style' : 'dashed'},
                    res=self.with_resources
                )
                graphviz_add_obj_edges(result,s_id,step_dict["parts"],**args)

                args["attr"] = {'color' : 'red'}
                graphviz_add_obj_edges(result,s_id,step_dict["tools"],**args)

                args["attr"] = {'color' : 'brown'}
                graphviz_add_obj_edges(result,s_id,step_dict["results"],**args)

        if self.with_resources:
            for ref in schedule.steps:
                s_id = 's_' + str(ref.step_nr)
                step = store.get_step(ref.step_id)

                for fil in step.files.values():
                    result.add_edge('r_' + fil.blob_id,s_id,color='orange')
                for img in step.images.values():
                    result.add_edge('r_' + img.blob_id,s_id,color='green')

        result.draw(path,prog='dot')

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This is free and unencumbered software released into the public domain.
#
# Anyone is free to copy, modify, publish, use, compile, sell, or
# distribute this software, either in source code form or as a compiled
# binary, for any purpose, commercial or non-commercial, and by any
# means.
#
# In jurisdictions that recognize copyright laws, the author or authors
# of this software dedicate any and all copyright interest in the
# software to the public domain. We make this dedication for the benefit
# of the public at large and to the detriment of our heirs and
# successors. We intend this dedication to be an overt act of
# relinquishment in perpetuity of all present and future rights to this
# software under copyright law.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# For more information, please refer to <http://unlicense.org/>
#
"""DVCS View

This tool helps to get an overview of the status of Git and Mercurial
repositories. The script searches for all repos in your workspaces and prints
a short status overview. It checks for uncommited changes in the working
directory and if configured pull/push-repos are in sync.

[DVCS View](http://github.com/samuelspiza/dvcsview) is hosted on Github.

The [template](http://gist.github.com/258034) contains examples for the
configuration of DVCS View.
"""

__author__ = "Samuel Spiza <sam.spiza@gmail.com>"
__version__ = "0.2"

import re
import os
import ConfigParser
import optparse
from subprocess import call, Popen, PIPE
import sys

CONFIG_FILENAMES = [os.path.expanduser("~/.dvcsview.conf"), "dvcsview.conf",
                    os.path.expanduser("~/dvcsview.ini"), "dvcsview.ini"]
SETTINGS = "settings"
WORKSPACES = "workspaces"
REPOS = "repos"
ALIAS = "alias"

def getOptions(argv):
    parser = optparse.OptionParser()
    parser.add_option("-t", "--targets",
                      dest="targets", metavar="HOSTS", default="",
                      help="Pull/push-repos on these hosts will be checked if "
                           "they are in sync. Argument is a comma separated "
                           "list of the hosts. Hosts can be IPv4-addresses or "
                           "domains. Windows volume letters (e.g. 'c:') and "
                           "1st level folder in the root directory "
                           "(e.g. '/home') in Linux systems are possible too. "
                           "Alias can be configured in '.dvcsview.conf'.")
    parser.add_option("-q", "--quiet",
                      dest="quiet", action="store_true", default=False,
                      help="Only one line per repository.")
    return parser.parse_args(argv)[0]

def getWorkspaces(config):
    workspaces = [os.path.expanduser(w[1]) for w in config.items(WORKSPACES)]
    for w in workspaces[:]:
        if not os.path.exists(w):
            print "ERROR: Workspace '%s' does not exist.\n" % w
            workspaces.remove(w)
    return workspaces

def findRepos(path, repos, targets=[], skip=[], quiet=False):
    entries = addRepo(path, repos, targets=targets, skip=skip, quiet=quiet)
    for entry in entries:
        newpath = os.path.join(path, entry)
        if os.path.isdir(newpath):
            findRepos(newpath, repos, targets=targets, skip=skip, quiet=quiet)

def addSingleRepo(path, repos, targets=[], skip=[], quiet=False):
    if not os.path.exists(path):
        print "ERROR: Repository '%s' does not exist.\n" % path
        return
    for repo in repos:
        if path == repo.path:
            print "ERROR: Repository '%s' is in a workspace.\n" % path
            return
    addRepo(path, repos, targets=targets, skip=skip, quiet=quiet)

def addRepo(path, repos, targets=[], skip=[], quiet=False):
    entries = os.listdir(path)
    if ".git" in entries and not "git" in skip:
        repos.append(Git(path, targets=targets, quiet=quiet))
        entries.remove(".git")
    elif ".hg" in entries and not "hg" in skip:
        repos.append(Hg(path, targets=targets, quiet=quiet))
        entries.remove(".hg")
    return entries

class WrappedFile:
    def __init__(self, path):
        self.file = open(path, "r")

    def readline(self):
        return self.file.readline().lstrip()

class Repo:
    def __init__(self, path, targets=[], quiet=False):
        self.path = path
        os.chdir(self.path)
        self.targets = targets
        self.quiet = quiet
        self.config = self.getConfig()
        ipseg = "[12]?[0-9]{1,2}"
        ip = "(?<=@)(?:" + ipseg + "\.){3}" + ipseg + "(?=[:/])"
        url = "(?<=[/@\.])[a-z0-9-]*\.(?:com|de|net|org)(?=[:/])"
        d = "^(?:/[a-z]+(?=/)|[a-z]:)"
        regexp = "|".join([ip, url, d])
        self.re = re.compile(regexp)
        self.fetch()
        self.statusstring = self.buildStatusString()

    def getConfig(self):
        config = None
        if os.path.exists(self.configFile):
            config = ConfigParser.ConfigParser()
            config.readfp(open(self.configFile))
        return config

    def buildStatusString(self):
        status = self.getStatus()
        if self.isClean(status):
            isClean = "CLEAN    "
        else:
            isClean = "NOT CLEAN"
        text = self.path
        if not self.quiet:
            text = "\n  ".join([text] + self.getWarnings(status))
        return "%s %s %s" % (isClean, self.typ, text)

    def pipe(self, command):
        pipe = Popen(command, shell=True, stdout=PIPE).stdout
        return [l.strip() for l in pipe.readlines()]

    def getWarnings(self, status):
        count = [[c, 0] for c in self.count]
        mod = []
        for line in status:
            if line not in self.skip:
                for rep in self.replace:
                    line = re.sub(rep[0], rep[1], line)
                a = False
                for i in range(len(count)):
                    m = re.search(count[i][0][0], line)
                    if m is not None:
                        count[i][1] += 1
                        a = True
                if not a:
                    mod.append(line)
        for c in count:
            if 0 < c[1]:
                if 1 == c[1]:
                    mod[0:0] = ["# %s: 1 file" % c[0][1]]
                else:
                    mod[0:0] = ["# %s: %s files" % (c[0][1], c[1])]
        return mod

class Git(Repo):
    typ = "Git"
    configFile = ".git/config"
    count = []
    count.append(["^#\t(?![a-z ]*:)","untracked"])
    count.append(["^#\tdeleted:", "deleted"])
    count.append(["^#\trenamed:", "renamed"])
    count.append(["^#\tmodified:", "modified"])
    count.append(["^#\tnew file:", "new file"])
    replace = [("Your branch is ahead of '\w*/\w*' by ([0-9]* commits?).",
                "ahead: \g<1>"),
               ("Changes to be committed:", "Changes to be committed")]
    skip = ["# Changed but not updated:", "# Untracked files:"]

    def __init__(self, path, targets=[], quiet=False):
        self.trackingbranches = None
        Repo.__init__(self, path, targets=targets, quiet=quiet)

    def getConfig(self):
        config = ConfigParser.ConfigParser()
        config.readfp(WrappedFile(self.configFile))
        return config

    def fetch(self):
        for remote in self.getRemotes():
            print "fetch %s" % self.path
            call("git fetch %s" % remote, shell=True)

    def getRemotes(self):
        remotesections = [s for s in self.config.sections()
                          if s.startswith("remote")]
        remotes = []
        for r in remotesections:
            m = self.re.search(self.config.get(r, "url"))
            if m.group(0) in self.targets:
                remotes.append(r[8:-1])
        return remotes

    def getStatus(self):
        if self.trackingbranches is None:
            self.trackingbranches = self.getTrackingBranches()
        status = self.pipe("git status")
        self.trackingbranches[status[0][12:]] = status
        return status

    def getTrackingBranches(self):
        branches = [s for s in self.config.sections()
                      if s.startswith("branch")]
        return dict([(b[8:-1], None) for b in branches
                     if self.config.has_option(b, "merge")])

    def isClean(self, status):
        b = status[-1][:17] == "nothing to commit"
        if b:
            self.allBranches(status)
        return b

    def allBranches(self, status):
        defaultbranch = status[0][12:]
        for b in self.trackingbranches.items():
            if b[0] != status[0][12:]:
                call("git checkout %s" % b[0], shell=True)
                status = self.getStatus()
        if status[0][12:] != defaultbranch:
            call("git checkout %s" % defaultbranch, shell=True)

    def getWarnings(self, status):
        warnings = []
        for status in self.trackingbranches.values():
            if status is not None:
                mod = Repo.getWarnings(self, status[1:])
                warn = ["  %s" % s for s in mod
                        if s.startswith("# ") and not s[1:].startswith("  ")]
                if 0 < len(warn):
                    warn[0:0] = ["# %s" % status[0][5:]]
                    warnings.extend(warn)
        return warnings

class Hg(Repo):
    typ = "Hg "
    configFile = ".hg/hgrc"
    count = [("^\?", "?"), ("^M", "M")]
    replace = []
    skip = []

    def fetch(self):
        paths = []
        if self.config is not None and self.config.has_option("paths", 'default'):
            paths.append(('default', "hg incoming"))
            if self.config.has_option("paths", 'default-push'):
                paths.append(('default-push', "hg outgoing"))
            else:
                paths.append(('default', "hg outgoing"))
        self.inout = []
        for p in paths:
            m = self.re.search(self.config.get("paths", p[0]))
            if m.group(0) in self.targets:
                line = self.traffic(p[1])
                if line is not None:
                    self.inout.append(line)

    def traffic(self, cmd):
        s = len([l for l in self.pipe(cmd) if l.startswith("changeset")])
        if 0 < s:
            return "# %s: %s changesets" % (cmd[3:], s)
        else:
            return None

    def getStatus(self):
        return self.pipe("hg status")

    def isClean(self, status):
        return 0 == len(status)

    def getWarnings(self, status):
        return Repo.getWarnings(self, status) + self.inout

def main(argv):
    config = ConfigParser.ConfigParser()
    config.read(CONFIG_FILENAMES)

    options = getOptions(argv)

    targets = options.targets
    # replace fetch alias with configured hosts
    if config.has_section(ALIAS):
        for opt in config.options(ALIAS):
            if targets == opt:
                targets = config.get(ALIAS, opt)
    # split comma separated list and strip elements
    targets = [t.strip() for t in targets.split(',') if 0 < len(t.strip())]

    # Parse the VCSs that shall be skipped.
    skip = []
    if config.has_section(SETTINGS):
        for opt in config.options(SETTINGS):
            if opt.startswith("skip.") and config.getboolean(SETTINGS, opt):
                skip.append(opt[5:])

    repos = []

    if config.has_section(WORKSPACES):
        workspaces = getWorkspaces(config)
        for workspace in workspaces:
            # creates 'Git' or 'Hg' objects and appends them to 'repos'
            findRepos(workspace, repos, targets=targets, skip=skip,
                      quiet=options.quiet)

    if config.has_section(REPOS):
        singlerepos = config.items(REPOS)
        for path in singlerepos:
            # creates 'Git' or 'Hg' objects and appends them to 'repos'
            addSingleRepo(path[1], repos, targets=targets, skip=skip,
                          quiet=options.quiet)

    for repo in repos:
        print repo.statusstring

    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
<commit_msg>Improved output of fetching changesets from remote repos<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This is free and unencumbered software released into the public domain.
#
# Anyone is free to copy, modify, publish, use, compile, sell, or
# distribute this software, either in source code form or as a compiled
# binary, for any purpose, commercial or non-commercial, and by any
# means.
#
# In jurisdictions that recognize copyright laws, the author or authors
# of this software dedicate any and all copyright interest in the
# software to the public domain. We make this dedication for the benefit
# of the public at large and to the detriment of our heirs and
# successors. We intend this dedication to be an overt act of
# relinquishment in perpetuity of all present and future rights to this
# software under copyright law.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
# OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# For more information, please refer to <http://unlicense.org/>
#
"""DVCS View

This tool helps to get an overview of the status of Git and Mercurial
repositories. The script searches for all repos in your workspaces and prints
a short status overview. It checks for uncommited changes in the working
directory and if configured pull/push-repos are in sync.

[DVCS View](http://github.com/samuelspiza/dvcsview) is hosted on Github.

The [template](http://gist.github.com/258034) contains examples for the
configuration of DVCS View.
"""

__author__ = "Samuel Spiza <sam.spiza@gmail.com>"
__version__ = "0.2.1"

import re
import os
import ConfigParser
import optparse
from subprocess import call, Popen, PIPE
import sys

CONFIG_FILENAMES = [os.path.expanduser("~/.dvcsview.conf"), "dvcsview.conf",
                    os.path.expanduser("~/dvcsview.ini"), "dvcsview.ini"]
SETTINGS = "settings"
WORKSPACES = "workspaces"
REPOS = "repos"
ALIAS = "alias"

def getOptions(argv):
    parser = optparse.OptionParser()
    parser.add_option("-t", "--targets",
                      dest="targets", metavar="HOSTS", default="",
                      help="Pull/push-repos on these hosts will be checked if "
                           "they are in sync. Argument is a comma separated "
                           "list of the hosts. Hosts can be IPv4-addresses or "
                           "domains. Windows volume letters (e.g. 'c:') and "
                           "1st level folder in the root directory "
                           "(e.g. '/home') in Linux systems are possible too. "
                           "Alias can be configured in '.dvcsview.conf'.")
    parser.add_option("-q", "--quiet",
                      dest="quiet", action="store_true", default=False,
                      help="Only one line per repository.")
    return parser.parse_args(argv)[0]

def getWorkspaces(config):
    workspaces = [os.path.expanduser(w[1]) for w in config.items(WORKSPACES)]
    for w in workspaces[:]:
        if not os.path.exists(w):
            print "ERROR: Workspace '%s' does not exist.\n" % w
            workspaces.remove(w)
    return workspaces

def findRepos(path, repos, targets=[], skip=[], quiet=False):
    entries = addRepo(path, repos, targets=targets, skip=skip, quiet=quiet)
    for entry in entries:
        newpath = os.path.join(path, entry)
        if os.path.isdir(newpath):
            findRepos(newpath, repos, targets=targets, skip=skip, quiet=quiet)

def addSingleRepo(path, repos, targets=[], skip=[], quiet=False):
    if not os.path.exists(path):
        print "ERROR: Repository '%s' does not exist.\n" % path
        return
    for repo in repos:
        if path == repo.path:
            print "ERROR: Repository '%s' is in a workspace.\n" % path
            return
    addRepo(path, repos, targets=targets, skip=skip, quiet=quiet)

def addRepo(path, repos, targets=[], skip=[], quiet=False):
    entries = os.listdir(path)
    if ".git" in entries and not "git" in skip:
        repos.append(Git(path, targets=targets, quiet=quiet))
        entries.remove(".git")
    elif ".hg" in entries and not "hg" in skip:
        repos.append(Hg(path, targets=targets, quiet=quiet))
        entries.remove(".hg")
    return entries

class WrappedFile:
    def __init__(self, path):
        self.file = open(path, "r")

    def readline(self):
        return self.file.readline().lstrip()

class Repo:
    def __init__(self, path, targets=[], quiet=False):
        self.path = path
        os.chdir(self.path)
        self.targets = targets
        self.quiet = quiet
        self.config = self.getConfig()
        ipseg = "[12]?[0-9]{1,2}"
        ip = "(?<=@)(?:" + ipseg + "\.){3}" + ipseg + "(?=[:/])"
        url = "(?<=[/@\.])[a-z0-9-]*\.(?:com|de|net|org)(?=[:/])"
        d = "^(?:/[a-z]+(?=/)|[a-z]:)"
        regexp = "|".join([ip, url, d])
        self.re = re.compile(regexp)
        self.fetch()
        self.statusstring = self.buildStatusString()

    def getConfig(self):
        config = None
        if os.path.exists(self.configFile):
            config = ConfigParser.ConfigParser()
            config.readfp(open(self.configFile))
        return config

    def fetch(self):
        for r in self.getRemotes():
            m = self.re.search(r['url'])
            if m is not None and m.group(0) in self.targets:
                print "fetch %s (%s)" % (r['url'], r['comment'])
                self.fetchRemote(r)

    def buildStatusString(self):
        status = self.getStatus()
        if self.isClean(status):
            isClean = "CLEAN    "
        else:
            isClean = "NOT CLEAN"
        text = self.path
        if not self.quiet:
            text = "\n  ".join([text] + self.getWarnings(status))
        return "%s %s %s" % (isClean, self.typ, text)

    def pipe(self, command):
        pipe = Popen(command, shell=True, stdout=PIPE).stdout
        return [l.strip() for l in pipe.readlines()]

    def getWarnings(self, status):
        count = [[c, 0] for c in self.count]
        mod = []
        for line in status:
            if line not in self.skip:
                for rep in self.replace:
                    line = re.sub(rep[0], rep[1], line)
                a = False
                for i in range(len(count)):
                    m = re.search(count[i][0][0], line)
                    if m is not None:
                        count[i][1] += 1
                        a = True
                if not a:
                    mod.append(line)
        for c in count:
            if 0 < c[1]:
                if 1 == c[1]:
                    mod[0:0] = ["# %s: 1 file" % c[0][1]]
                else:
                    mod[0:0] = ["# %s: %s files" % (c[0][1], c[1])]
        return mod

class Git(Repo):
    typ = "Git"
    configFile = ".git/config"
    count = []
    count.append(["^#\t(?![a-z ]*:)","untracked"])
    count.append(["^#\tdeleted:", "deleted"])
    count.append(["^#\trenamed:", "renamed"])
    count.append(["^#\tmodified:", "modified"])
    count.append(["^#\tnew file:", "new file"])
    replace = [("Your branch is ahead of '\w*/\w*' by ([0-9]* commits?).",
                "ahead: \g<1>"),
               ("Changes to be committed:", "Changes to be committed")]
    skip = ["# Changed but not updated:", "# Untracked files:"]

    def __init__(self, path, targets=[], quiet=False):
        self.trackingbranches = None
        Repo.__init__(self, path, targets=targets, quiet=quiet)

    def getConfig(self):
        config = ConfigParser.ConfigParser()
        config.readfp(WrappedFile(self.configFile))
        return config

    def fetchRemote(self, r):
        call(r['cmd'], shell=True)

    def getRemotes(self):
        remotesections = [s for s in self.config.sections()
                          if s.startswith("remote")]
        remotes = []
        for r in remotesections:
            remotes.append({'cmd': "git fetch " + r[8:-1],
                            'url': self.config.get(r, "url"),
                            'comment': r[8:-1]})
        return remotes

    def getStatus(self):
        if self.trackingbranches is None:
            self.trackingbranches = self.getTrackingBranches()
        status = self.pipe("git status")
        self.trackingbranches[status[0][12:]] = status
        return status

    def getTrackingBranches(self):
        branches = [s for s in self.config.sections()
                      if s.startswith("branch")]
        return dict([(b[8:-1], None) for b in branches
                     if self.config.has_option(b, "merge")])

    def isClean(self, status):
        b = status[-1][:17] == "nothing to commit"
        if b:
            self.allBranches(status)
        return b

    def allBranches(self, status):
        defaultbranch = status[0][12:]
        for b in self.trackingbranches.items():
            if b[0] != status[0][12:]:
                call("git checkout %s" % b[0], shell=True)
                status = self.getStatus()
        if status[0][12:] != defaultbranch:
            call("git checkout %s" % defaultbranch, shell=True)

    def getWarnings(self, status):
        warnings = []
        for status in self.trackingbranches.values():
            if status is not None:
                mod = Repo.getWarnings(self, status[1:])
                warn = ["  %s" % s for s in mod
                        if s.startswith("# ") and not s[1:].startswith("  ")]
                if 0 < len(warn):
                    warn[0:0] = ["# %s" % status[0][5:]]
                    warnings.extend(warn)
        return warnings

class Hg(Repo):
    typ = "Hg "
    configFile = ".hg/hgrc"
    count = [("^\?", "?"), ("^M", "M")]
    replace = []
    skip = []

    def __init__(self, path, targets=[], quiet=False):
        self.inout = []
        Repo.__init__(self, path, targets=targets, quiet=quiet)

    def getRemotes(self):
        remotes = []
        if self.config is not None and self.config.has_option("paths", 'default'):
            default = self.config.get("paths", "default")
            remotes.extend([{'cmd': "hg incoming",
                             'url': default,
                             'comment': "incoming"},
                            {'cmd': "hg outgoing",
                             'url': default,
                             'comment': "outgoing"}])
            if self.config.has_option("paths", 'default-push'):
                remotes[1]['url'] = self.config.get("paths", "default-push")
        return remotes

    def fetchRemote(self, r):
        s = len([l for l in self.pipe(r['cmd']) if l.startswith("changeset")])
        if 0 < s:
            self.inout.append("# %s: %s changesets" % (cmd[3:], s))

    def getStatus(self):
        return self.pipe("hg status")

    def isClean(self, status):
        return 0 == len(status)

    def getWarnings(self, status):
        return Repo.getWarnings(self, status) + self.inout

def main(argv):
    config = ConfigParser.ConfigParser()
    config.read(CONFIG_FILENAMES)

    options = getOptions(argv)

    targets = options.targets
    # replace fetch alias with configured hosts
    if config.has_section(ALIAS):
        for opt in config.options(ALIAS):
            if targets == opt:
                targets = config.get(ALIAS, opt)
    # split comma separated list and strip elements
    targets = [t.strip() for t in targets.split(',') if 0 < len(t.strip())]

    # Parse the VCSs that shall be skipped.
    skip = []
    if config.has_section(SETTINGS):
        for opt in config.options(SETTINGS):
            if opt.startswith("skip.") and config.getboolean(SETTINGS, opt):
                skip.append(opt[5:])

    repos = []

    if config.has_section(WORKSPACES):
        workspaces = getWorkspaces(config)
        for workspace in workspaces:
            # creates 'Git' or 'Hg' objects and appends them to 'repos'
            findRepos(workspace, repos, targets=targets, skip=skip,
                      quiet=options.quiet)

    if config.has_section(REPOS):
        singlerepos = config.items(REPOS)
        for path in singlerepos:
            # creates 'Git' or 'Hg' objects and appends them to 'repos'
            addSingleRepo(path[1], repos, targets=targets, skip=skip,
                          quiet=options.quiet)

    for repo in repos:
        print repo.statusstring

    return 0

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


<commit_before>import json
import logging
from optparse import make_option
import datetime

import requests
from celery_tasktree import task_with_callbacks, TaskTree

from django.conf import settings as django_settings
from django.core.management.base import BaseCommand, CommandError

from amo.utils import timestamp_index

from addons.cron import reindex_addons, reindex_apps
from apps.addons.search import setup_mapping as put_amo_mapping
from bandwagon.cron import reindex_collections
from compat.cron import compatibility_report
from stats.cron import index_latest_stats
from stats.search import setup_indexes as put_stats_mapping
from users.cron import reindex_users

from lib.es.models import Reindexing

if django_settings.MARKETPLACE:
    from mkt.stats.cron import index_latest_mkt_stats, index_mkt_stats
    from mkt.stats.search import setup_mkt_indexes as put_mkt_stats_mapping

_INDEXES = {'stats': [index_latest_stats],
            'apps': [reindex_addons,
                     reindex_apps,
                     reindex_collections,
                     reindex_users]}

if django_settings.MARKETPLACE:
    _INDEXES['stats'].extend([index_mkt_stats, index_latest_mkt_stats])
else:
    _INDEXES['apps'].append(compatibility_report)

logger = logging.getLogger('z.elasticsearch')
DEFAULT_NUM_REPLICAS = 0
DEFAULT_NUM_SHARDS = 3

if hasattr(django_settings, 'ES_HOSTS'):
    base_url = django_settings.ES_HOSTS[0]
else:
    base_url = 'http://127.0.0.1:9200'

def url(path):
    return '%s%s' % (base_url, path)

def _action(name, **kw):
    return {name: kw}

def call_es(path, *args, **kw):
    method = kw.pop('method', 'GET')
    status = kw.pop('status', 200)
    if isinstance(status, int):
        status = [status]

    if not path.startswith('/'):
        path = '/' + path

    method = getattr(requests, method.lower())
    res = method(url(path), *args, **kw)

    if res.status_code not in status:
        raise CommandError(res.content)
    return res

@task_with_callbacks
def delete_indexes(indexes):
    """Removes the indexes.

    - indexes: list of indexes names to remove.
    """
    # now call the server - can we do this with a single call?
    for index in indexes:
        call_es(index, method='DELETE')

@task_with_callbacks
def run_aliases_actions(actions):
    """Run actions on aliases.

     - action: list of action/index/alias items
    """
    # we also want to rename or delete the current index in case we have one
    dump = []

    for action, index, alias in actions:
        dump.append({action: {'index': index, 'alias': alias}})

    # now call the server
    call_es('_aliases', json.dumps(dict(actions=dump)), method='POST')

@task_with_callbacks
def create_mapping(new_index, alias, num_replicas=DEFAULT_NUM_REPLICAS,
                  num_shards=DEFAULT_NUM_SHARDS):
    """Creates a mapping for the new index.

    - new_index: new index name.
    - alias: alias name
    - num_replicas: number of replicas in ES
    - num_shards: number of shards in ES
    """
    if requests.head(url('/' + alias)).status_code == 200:
        res = call_es('%s/_settings' % (alias)).json
        idx_settings = res.get(alias, {}).get('settings', {})
    else:
        idx_settings = {}

    settings = {
        'number_of_replicas': idx_settings.get('number_of_replicas',
                                                num_replicas),
        'number_of_shards': idx_settings.get('number_of_shards',
                                                num_shards)
    }

    # Create mapping.without aliases since we do it manually
    if not 'stats' in alias:
        put_amo_mapping(new_index, aliased=False)
    else:
        put_stats_mapping(new_index, aliased=False)
        if django_settings.MARKETPLACE:
            put_mkt_stats_mapping(new_index, aliased=False)

    # Create new index
    index_url = url('/%s' % new_index)

    # if the index already exists we can keep it
    if requests.head(index_url).status_code == 200:
        return

    call_es(index_url, json.dumps(settings), method='PUT',
            status=(200, 201))

@task_with_callbacks
def create_index(index, is_stats):
    """Create the index.

    - index: name of the index
    - is_stats: if True, we're indexing stats
    """
    indexers = is_stats and _INDEXES['stats'] or _INDEXES['apps']
    for indexer in indexers:
        indexer(index, aliased=False)

def database_flagged():
    """Returns True if the Database is being indexed"""
    return Reindexing.objects.exists()

@task_with_callbacks
def flag_database(new_index, old_index, alias):
    """Flags the database to indicate that the reindexing has started."""
    return Reindexing.objects.create(new_index=new_index, old_index=old_index,
                                     alias=alias,
                                     start_date=datetime.datetime.now())

@task_with_callbacks
def unflag_database():
    """Unflag the database to indicate that the reindexing is over."""
    Reindexing.objects.all().delete()

_SUMMARY = """
*** Reindexation done ***

Reindexed %d indexes.

Current Aliases configuration:

%s

"""

class Command(BaseCommand):
    help = 'Reindex all ES indexes'
    option_list = BaseCommand.option_list + (
        make_option('--prefix', action='store',
                    help='Indexes prefixes, like test_',
                    default=''),
        make_option('--force', action='store_true',
                    help=('Bypass the database flag that says '
                          'another indexation is ongoing'),
                    default=False),
        make_option('--wipe', action='store_true',
                    help=('Wipe ES from any content first. This option '
                          'will destroy anything that is in ES!'),
                    default=False),

    )

    def handle(self, *args, **kwargs):
        """Reindexing work.

        Creates a Tasktree that creates new indexes
        over the old ones so the search feature
        works while the indexation occurs
        """
        force = kwargs.get('force', False)

        if database_flagged() and not force:
            raise CommandError('Indexation already occuring - use --force to '
                               'bypass')

        prefix = kwargs.get('prefix', '')
        logger.info('Starting the reindexation')

        if kwargs.get('wipe', False):
            confirm = raw_input("Are you sure you want to wipe all data from "
                                "ES ? (yes/no): ")

            while confirm not in ('yes', 'no'):
                confirm = raw_input('Please enter either "yes" or "no": ')

            if confirm == 'yes':
                unflag_database()
                requests.delete(url('/'))
            else:
                raise CommandError("Aborted.")
        elif force:
            unflag_database()

        # Get list current aliases at /_aliases.
        all_aliases = requests.get(url('/_aliases')).json

        # building the list of indexes
        indexes = set([prefix + index for index in
                       django_settings.ES_INDEXES.values()])

        actions = []

        def add_action(*elmt):
            if elmt in actions:
                return
            actions.append(elmt)

        all_aliases = all_aliases.items()

        # creating a task tree
        tree = TaskTree()
        last_action = None

        to_remove = []

        # for each index, we create a new time-stamped index
        for alias in indexes:
            old_index = None

            for aliased_index, alias_ in all_aliases:
                if alias in alias_['aliases'].keys():
                    # mark the index to be removed later
                    old_index = aliased_index
                    to_remove.append(aliased_index)

                    # mark the alias to be removed as well
                    add_action('remove', aliased_index, alias)

            # create a new index, using the alias name with a timestamp
            new_index = timestamp_index(alias)

            # flag the database
            step1 = tree.add_task(flag_database, args=[new_index, old_index,
                                                       alias])
            step2 = step1.add_task(create_mapping, args=[new_index, alias])
            step3 = step2.add_task(create_index, args=[new_index,
                                                       'stats' in alias])
            last_action = step3

            # adding new index to the alias
            add_action('add', new_index, alias)

        # Alias the new index and remove the old aliases, if any.
        renaming_step = last_action.add_task(run_aliases_actions,
                                             args=[actions])

        # unflag the database - there's no need to duplicate the
        # indexing anymore
        delete = renaming_step.add_task(unflag_database)

        # Delete the old indexes, if any
        delete.add_task(delete_indexes, args=[to_remove])

        # let's do it
        tree.apply_async()

        # let's return the /_aliases values
        aliases = call_es('_aliases').json
        aliases = json.dumps(aliases, sort_keys=True, indent=4)
        return _SUMMARY % (len(indexes), aliases)
<commit_msg>ES_HOSTS elements come with no scheme<commit_after>import json
import logging
from optparse import make_option
import datetime

import requests
from celery_tasktree import task_with_callbacks, TaskTree

from django.conf import settings as django_settings
from django.core.management.base import BaseCommand, CommandError

from amo.utils import timestamp_index

from addons.cron import reindex_addons, reindex_apps
from apps.addons.search import setup_mapping as put_amo_mapping
from bandwagon.cron import reindex_collections
from compat.cron import compatibility_report
from stats.cron import index_latest_stats
from stats.search import setup_indexes as put_stats_mapping
from users.cron import reindex_users

from lib.es.models import Reindexing

if django_settings.MARKETPLACE:
    from mkt.stats.cron import index_latest_mkt_stats, index_mkt_stats
    from mkt.stats.search import setup_mkt_indexes as put_mkt_stats_mapping

_INDEXES = {'stats': [index_latest_stats],
            'apps': [reindex_addons,
                     reindex_apps,
                     reindex_collections,
                     reindex_users]}

if django_settings.MARKETPLACE:
    _INDEXES['stats'].extend([index_mkt_stats, index_latest_mkt_stats])
else:
    _INDEXES['apps'].append(compatibility_report)

logger = logging.getLogger('z.elasticsearch')
DEFAULT_NUM_REPLICAS = 0
DEFAULT_NUM_SHARDS = 3

if hasattr(django_settings, 'ES_HOSTS'):
    base_url = django_settings.ES_HOSTS[0]
    # The configuration comes with no scheme.
    base_url = 'http://%s' % base_url
else:
    base_url = 'http://127.0.0.1:9200'

def url(path):
    return '%s%s' % (base_url, path)

def _action(name, **kw):
    return {name: kw}

def call_es(path, *args, **kw):
    method = kw.pop('method', 'GET')
    status = kw.pop('status', 200)
    if isinstance(status, int):
        status = [status]

    if not path.startswith('/'):
        path = '/' + path

    method = getattr(requests, method.lower())
    res = method(url(path), *args, **kw)

    if res.status_code not in status:
        raise CommandError(res.content)
    return res

@task_with_callbacks
def delete_indexes(indexes):
    """Removes the indexes.

    - indexes: list of indexes names to remove.
    """
    # now call the server - can we do this with a single call?
    for index in indexes:
        call_es(index, method='DELETE')

@task_with_callbacks
def run_aliases_actions(actions):
    """Run actions on aliases.

     - action: list of action/index/alias items
    """
    # we also want to rename or delete the current index in case we have one
    dump = []

    for action, index, alias in actions:
        dump.append({action: {'index': index, 'alias': alias}})

    # now call the server
    call_es('_aliases', json.dumps(dict(actions=dump)), method='POST')

@task_with_callbacks
def create_mapping(new_index, alias, num_replicas=DEFAULT_NUM_REPLICAS,
                  num_shards=DEFAULT_NUM_SHARDS):
    """Creates a mapping for the new index.

    - new_index: new index name.
    - alias: alias name
    - num_replicas: number of replicas in ES
    - num_shards: number of shards in ES
    """
    if requests.head(url('/' + alias)).status_code == 200:
        res = call_es('%s/_settings' % (alias)).json
        idx_settings = res.get(alias, {}).get('settings', {})
    else:
        idx_settings = {}

    settings = {
        'number_of_replicas': idx_settings.get('number_of_replicas',
                                                num_replicas),
        'number_of_shards': idx_settings.get('number_of_shards',
                                                num_shards)
    }

    # Create mapping.without aliases since we do it manually
    if not 'stats' in alias:
        put_amo_mapping(new_index, aliased=False)
    else:
        put_stats_mapping(new_index, aliased=False)
        if django_settings.MARKETPLACE:
            put_mkt_stats_mapping(new_index, aliased=False)

    # Create new index
    index_url = url('/%s' % new_index)

    # if the index already exists we can keep it
    if requests.head(index_url).status_code == 200:
        return

    call_es(index_url, json.dumps(settings), method='PUT',
            status=(200, 201))

@task_with_callbacks
def create_index(index, is_stats):
    """Create the index.

    - index: name of the index
    - is_stats: if True, we're indexing stats
    """
    indexers = is_stats and _INDEXES['stats'] or _INDEXES['apps']
    for indexer in indexers:
        indexer(index, aliased=False)

def database_flagged():
    """Returns True if the Database is being indexed"""
    return Reindexing.objects.exists()

@task_with_callbacks
def flag_database(new_index, old_index, alias):
    """Flags the database to indicate that the reindexing has started."""
    return Reindexing.objects.create(new_index=new_index, old_index=old_index,
                                     alias=alias,
                                     start_date=datetime.datetime.now())

@task_with_callbacks
def unflag_database():
    """Unflag the database to indicate that the reindexing is over."""
    Reindexing.objects.all().delete()

_SUMMARY = """
*** Reindexation done ***

Reindexed %d indexes.

Current Aliases configuration:

%s

"""

class Command(BaseCommand):
    help = 'Reindex all ES indexes'
    option_list = BaseCommand.option_list + (
        make_option('--prefix', action='store',
                    help='Indexes prefixes, like test_',
                    default=''),
        make_option('--force', action='store_true',
                    help=('Bypass the database flag that says '
                          'another indexation is ongoing'),
                    default=False),
        make_option('--wipe', action='store_true',
                    help=('Wipe ES from any content first. This option '
                          'will destroy anything that is in ES!'),
                    default=False),

    )

    def handle(self, *args, **kwargs):
        """Reindexing work.

        Creates a Tasktree that creates new indexes
        over the old ones so the search feature
        works while the indexation occurs
        """
        force = kwargs.get('force', False)

        if database_flagged() and not force:
            raise CommandError('Indexation already occuring - use --force to '
                               'bypass')

        prefix = kwargs.get('prefix', '')
        logger.info('Starting the reindexation')

        if kwargs.get('wipe', False):
            confirm = raw_input("Are you sure you want to wipe all data from "
                                "ES ? (yes/no): ")

            while confirm not in ('yes', 'no'):
                confirm = raw_input('Please enter either "yes" or "no": ')

            if confirm == 'yes':
                unflag_database()
                requests.delete(url('/'))
            else:
                raise CommandError("Aborted.")
        elif force:
            unflag_database()

        # Get list current aliases at /_aliases.
        all_aliases = requests.get(url('/_aliases')).json

        # building the list of indexes
        indexes = set([prefix + index for index in
                       django_settings.ES_INDEXES.values()])

        actions = []

        def add_action(*elmt):
            if elmt in actions:
                return
            actions.append(elmt)

        all_aliases = all_aliases.items()

        # creating a task tree
        tree = TaskTree()
        last_action = None

        to_remove = []

        # for each index, we create a new time-stamped index
        for alias in indexes:
            old_index = None

            for aliased_index, alias_ in all_aliases:
                if alias in alias_['aliases'].keys():
                    # mark the index to be removed later
                    old_index = aliased_index
                    to_remove.append(aliased_index)

                    # mark the alias to be removed as well
                    add_action('remove', aliased_index, alias)

            # create a new index, using the alias name with a timestamp
            new_index = timestamp_index(alias)

            # flag the database
            step1 = tree.add_task(flag_database, args=[new_index, old_index,
                                                       alias])
            step2 = step1.add_task(create_mapping, args=[new_index, alias])
            step3 = step2.add_task(create_index, args=[new_index,
                                                       'stats' in alias])
            last_action = step3

            # adding new index to the alias
            add_action('add', new_index, alias)

        # Alias the new index and remove the old aliases, if any.
        renaming_step = last_action.add_task(run_aliases_actions,
                                             args=[actions])

        # unflag the database - there's no need to duplicate the
        # indexing anymore
        delete = renaming_step.add_task(unflag_database)

        # Delete the old indexes, if any
        delete.add_task(delete_indexes, args=[to_remove])

        # let's do it
        tree.apply_async()

        # let's return the /_aliases values
        aliases = call_es('_aliases').json
        aliases = json.dumps(aliases, sort_keys=True, indent=4)
        return _SUMMARY % (len(indexes), aliases)
<|endoftext|>


<commit_before># encoding: utf-8
"""
server.py

Created by David Farrar  on 2011-11-30.
Copyright (c) 2011-2013  Exa Networks. All rights reserved.
"""

import socket
import errno

from exaproxy.network.errno_list import errno_block

def ishex (s):
	return bool(s) and not bool(s.strip('0123456789abcdefABCDEF'))

def count_quotes (data):
	return data.count('"') - data.count('\\"')

class Client (object):
	eor = ['\r\n\r\n', '\n\n']

	def __init__(self, name, sock, peer, logger):
		self.name = name
		self.sock = sock
		self.peer = peer
		self.reader = self._read(sock)
		self.writer = self._write(sock)

		self.log = logger
		self.blockupload = None

		# start the _read coroutine
		self.reader.next()

	def checkRequest (self, r_buffer, seek=0):
		# XXX: max buffer size
		for eor in self.eor:
			pos = r_buffer[seek:].find(eor)
			if pos == -1: continue

			buff = r_buffer[:seek+pos]
			if not buff: continue

			if not count_quotes(buff) % 2:  # we have matching pairs
				return buff + eor, r_buffer[seek+pos+len(eor):], seek

			seek += pos + len(eor)

		return '', r_buffer, seek

	def checkChunkSize (self, r_buffer):
		# return a tuple : bool, length
		# * the bool is : is there more chunk to come
		# * the len contains the size of the chunk(s) extracted
		#   a size of None means that we could not decode as it is invalid

		total_len = 0

		while r_buffer:
			if not '\n' in r_buffer:
				if len(r_buffer) > 6:  # len('FFFF') + len(\r\n)
					return True, None
				return True, 0

			header,r_buffer = r_buffer.split('\n', 1)
			len_header = len(header) + 1

			if header.endswith('\r'):
				header = header[:-1]
				len_eol = 2
			else:
				len_eol = 1

			if ';' in header:
				header = header.split(';',1)[0]

			if not ishex(header):
				return True,None

			len_chunk = int(header, 16)

			# 0xFFFF is not enough - coad is complaining :p
			if len_chunk > 0x100000:
				return True,None

			if len_chunk == 0:
				total_len += len_header
				return False, total_len
			else:
				total = len_chunk + len_eol
				total_len += total + len_header
				r_buffer = r_buffer[total:]

		return True,total_len

	def _read (self, sock, read_size=64*1024):
		"""Coroutine managing data read from the client"""
		# yield request, content
		# request is the text that form the request header
		# content any text which is related to the current request after the headers

		yield ''

		r_buffer = ''
		nb_to_send = 0
		seek = 0
		processing = False

		# mode can be one of : request, chunk, extension, relay
		# request : we are reading the request (read all you can until a separator)
		# extra-headers : we are reading data until a separator
		# chunked : we are reading chunk-encoded darta
		# transfer : we are reading as much as requested in remaining
		# passthrough : read as much as can to be relayed

		mode = 'request'

		while True:
			try:
				while True:
					if not processing:
						data = sock.recv(read_size)
						if not data:
							break  # read failed so we abort
						self.log.debug("<< [%s]" % data.replace('\t','\\t').replace('\r','\\r').replace('\n','\\n'))
						r_buffer += data
					else:
						processing = False

					if mode == 'passthrough':
						yield '', r_buffer
						r_buffer = ''
						continue

					if mode in ('transfer','chunked') and nb_to_send:
						r_len = len(r_buffer)
						length = min(r_len, nb_to_send)

						# do not yield yet if we are chunked since the end of the chunk may
						# very well be in the rest of the data we just read
						# XXX: should be able to support pipelining by applying the same idea
						# when in transfer mode
						if mode == 'transfer' or r_len <= nb_to_send:
							_, extra_size = yield '', r_buffer[:length]

							r_buffer = r_buffer[length:]
							nb_to_send = nb_to_send - length + extra_size

							# we still have data to read before we can send more.
							if nb_to_send != 0:
								continue

							# we finished relaying, go back to reading requests
							elif mode == 'transfer':
								mode = 'request'

					if mode == 'chunked':
						# sum of the sizes of all chunks in our buffer
						chunked, new_to_send = self.checkChunkSize(r_buffer[nb_to_send:])

						if new_to_send is None:
							# could not read any chunk (data is invalid)
							break

						nb_to_send += new_to_send
						if chunked:
							# process the chunk in the if above
							processing = True
							continue

						mode = 'end-chunk'

					# seek is only set if we already passed once and found we needed more data to check
					if mode == 'end-chunk':
						if r_buffer[nb_to_send:].startswith('\r\n'):
							nb_to_send += 2
							processing = True
							mode = 'transfer'
							continue
						elif r_buffer[nb_to_send:].startswith('\n'):
							nb_to_send += 1
							processing = True
							mode = 'transfer'
							continue

						if not r_buffer[nb_to_send:]:
							yield '',''
							continue

						mode = 'extra-headers'
						seek = nb_to_send

					if mode == 'extra-headers':
						# seek is up to where we know there is no double CRLF
						related, r_buffer, seek = self.checkRequest(r_buffer,seek)
						yield '', related

						if not related:
							continue

						seek = 0
						mode = 'transfer'

					if mode != 'request':
						self.log.error('The programmers are monkeys - please give them bananas ..')
						self.log.error('the mode was spelled : [%s]' % mode)
						self.log.error('.. if it works, we are lucky - but it may work.')
						mode = 'request'

					# ignore EOL
					r_buffer = r_buffer.lstrip('\r\n')

					# check to see if we have read an entire request
					request, r_buffer, seek = self.checkRequest(r_buffer, seek)

					if not request:
						yield '', ''
						continue
					seek = 0
					processing = True

					# nb_to_send is how much we expect to need to get the rest of the request
					mode, nb_to_send = yield request, ''

				# break out of the outer loop as soon as we leave the inner loop
				# through normal execution
				break

			except socket.error, e:
				if e.args[0] in errno_block:
					yield '', ''
				else:
					break

		yield None,None

	def setPeer (self, peer):
		"""Set the claimed ip address for this client.
		Does not effect the ip address we try sending data to."""
		self.peer = peer

	def readData(self):
		request,content = self.reader.send(('transfer',0))
		return self.name, self.peer, request, content

	def readRelated(self, mode, remaining):
		request, content = self.reader.send((mode,remaining))
		return self.name, self.peer, request, content

	def _write(self, sock):
		"""Coroutine managing data sent to the client"""

		w_buffer = ''
		filename = yield None

		# check to see if we are returning data directly from a local file
		if filename is not None:
			try:
				# XXX: we must read from the file on demand rather than doing this
				with open(filename) as fd:
					w_buffer = fd.read()

				found = True, False, 0
			except IOError:
				found = None

			data = yield found
			w_buffer = data + w_buffer
		else:
			found = None

		data = yield found
		finished = False

		while True:
			try:
				while True:
					had_buffer = bool(w_buffer)

					if data is not None:
						w_buffer += data
					else:
						# We've finished downloading, even if the client hasn't yet
						finished = True

					if finished:
						if not w_buffer:
							break      # terminate the client connection
						elif data:
							self.log.error('Tried to send data to client after we told it to close. Dropping it.')

					if not had_buffer or data == '':
						sent = sock.send(w_buffer)
						if sent:
							self.log.debug(">> [%s]" % w_buffer[:sent].replace('\t','\\t').replace('\r','\\r').replace('\n','\\n'))
						w_buffer = w_buffer[sent:]
					else:
						sent = 0

					buffered = bool(w_buffer) or finished
					data = yield buffered, had_buffer, sent

				# break out of the outer loop as soon as we leave the inner loop
				# through normal execution
				yield None
				break

			except socket.error, e:
				if e.args[0] in errno_block:
					#self.log.info('interrupted when trying to sent %d bytes, will retry' % len(data))
					#self.log.info('reason: errno %d: %s' % (e.args[0], errno.errorcode.get(e.args[0], '<no errno name>')))
					data = yield bool(w_buffer) or finished, had_buffer, 0
				else:
					#self.log.critical('unexpected error writing on socket')
					#self.log.critical('reason, errno %d: %s' % (e.args[0], errno.errorcode.get(e.args[0], '<no errno name>')))
					yield None  # stop the client connection
					break  # and don't come back

		yield None

	def writeData(self, data):
		res = self.writer.send(data)
		return res

	def startData(self, command, data):
		# start the _write coroutine
		self.writer.next()

		if command == 'stream':
			self.writer.send(None)  # no local file
			res = self.writer.send(data)

		elif command == 'close':
			self.writer.send(None)  # no local file
			self.writer.send(data)
			res = self.writer.send(None)  # close the connection once the buffer is empty

		elif command == 'file':
			header, filename = data
			res = self.writer.send(filename)  # use local file
			self.writer.send(header)  # write the response headers before the file

			self.writer.send(None)  # close the connection once the buffer is empty
		else:
			res = None

		# buffered, had_buffer
		return res

	def restartData(self, command, data):
		self.writer = self._write(self.sock)
		return self.startData(command, data)

	def shutdown(self):
		try:
			self.sock.shutdown(socket.SHUT_RDWR)
			self.sock.close()
		except socket.error:
			pass

		self.writer.close()
		self.reader.close()
<commit_msg>more readable<commit_after># encoding: utf-8
"""
server.py

Created by David Farrar  on 2011-11-30.
Copyright (c) 2011-2013  Exa Networks. All rights reserved.
"""

import socket
import errno

from exaproxy.network.errno_list import errno_block

def ishex (s):
	return bool(s) and not bool(s.strip('0123456789abcdefABCDEF'))

def count_quotes (data):
	return data.count('"') - data.count('\\"')

class Client (object):
	eor = ['\r\n\r\n', '\n\n']

	def __init__(self, name, sock, peer, logger):
		self.name = name
		self.sock = sock
		self.peer = peer
		self.reader = self._read(sock)
		self.writer = self._write(sock)

		self.log = logger
		self.blockupload = None

		# start the _read coroutine
		self.reader.next()

	def checkRequest (self, r_buffer, seek=0):
		# XXX: max buffer size
		for eor in self.eor:
			pos = r_buffer[seek:].find(eor)
			if pos == -1: continue

			buff = r_buffer[:seek+pos]
			if not buff: continue

			if not count_quotes(buff) % 2:  # we have matching pairs
				return buff + eor, r_buffer[seek+pos+len(eor):], seek

			seek += pos + len(eor)

		return '', r_buffer, seek

	def checkChunkSize (self, r_buffer):
		# return a tuple : bool, length
		# * the bool is : is there more chunk to come
		# * the len contains the size of the chunk(s) extracted
		#   a size of None means that we could not decode as it is invalid

		total_len = 0

		while r_buffer:
			if not '\n' in r_buffer:
				if len(r_buffer) > 6:  # len('FFFF') + len(\r\n)
					return True, None
				return True, 0

			header,r_buffer = r_buffer.split('\n', 1)
			len_header = len(header) + 1

			if header.endswith('\r'):
				header = header[:-1]
				len_eol = 2
			else:
				len_eol = 1

			if ';' in header:
				header = header.split(';',1)[0]

			if not ishex(header):
				return True,None

			len_chunk = int(header, 16)

			# 0xFFFF is not enough - coad is complaining :p
			if len_chunk > 0x100000:
				return True,None

			if len_chunk == 0:
				total_len += len_header
				return False, total_len
			else:
				total = len_chunk + len_eol
				total_len += total + len_header
				r_buffer = r_buffer[total:]

		return True,total_len

	def _read (self, sock, read_size=64*1024):
		"""Coroutine managing data read from the client"""
		# yield request, content
		# request is the text that form the request header
		# content any text which is related to the current request after the headers

		yield ''

		r_buffer = ''
		nb_to_send = 0
		seek = 0
		processing = False

		# mode can be one of : request, chunk, extension, relay
		# request : we are reading the request (read all you can until a separator)
		# extra-headers : we are reading data until a separator
		# chunked : we are reading chunk-encoded darta
		# transfer : we are reading as much as requested in remaining
		# passthrough : read as much as can to be relayed

		mode = 'request'

		while True:
			try:
				while True:
					if not processing:
						data = sock.recv(read_size)
						if not data:
							break  # read failed so we abort
						self.log.debug("<< [%s]" % data.replace('\t','\\t').replace('\r','\\r').replace('\n','\\n'))
						r_buffer += data
					else:
						processing = False

					if mode == 'passthrough':
						yield '', r_buffer
						r_buffer = ''
						continue

					if nb_to_send:
						if mode == 'transfer' :
							r_len = len(r_buffer)
							length = min(r_len, nb_to_send)

							# do not yield yet if we are chunked since the end of the chunk may
							# very well be in the rest of the data we just read
							_, extra_size = yield '', r_buffer[:length]

							r_buffer = r_buffer[length:]
							nb_to_send = nb_to_send - length + extra_size

							# we still have data to read before we can send more.
							if nb_to_send != 0:
								continue
							mode = 'request'

						if mode == 'chunked':
							r_len = len(r_buffer)
							length = min(r_len, nb_to_send)

							# do not yield yet if we are chunked since the end of the chunk may
							# very well be in the rest of the data we just read
							if r_len <= nb_to_send:
								_, extra_size = yield '', r_buffer[:length]

								r_buffer = r_buffer[length:]
								nb_to_send = nb_to_send - length + extra_size

								# we still have data to read before we can send more.
								if nb_to_send != 0:
									continue

					if mode == 'chunked':
						# sum of the sizes of all chunks in our buffer
						chunked, new_to_send = self.checkChunkSize(r_buffer[nb_to_send:])

						if new_to_send is None:
							# could not read any chunk (data is invalid)
							break

						nb_to_send += new_to_send
						if chunked:
							# process the chunk in the if above
							processing = True
							continue

						mode = 'end-chunk'

					# seek is only set if we already passed once and found we needed more data to check
					if mode == 'end-chunk':
						if r_buffer[nb_to_send:].startswith('\r\n'):
							nb_to_send += 2
							processing = True
							mode = 'transfer'
							continue
						elif r_buffer[nb_to_send:].startswith('\n'):
							nb_to_send += 1
							processing = True
							mode = 'transfer'
							continue

						if not r_buffer[nb_to_send:]:
							yield '',''
							continue

						mode = 'extra-headers'
						seek = nb_to_send

					if mode == 'extra-headers':
						# seek is up to where we know there is no double CRLF
						related, r_buffer, seek = self.checkRequest(r_buffer,seek)
						yield '', related

						if not related:
							continue

						seek = 0
						mode = 'transfer'

					if mode != 'request':
						self.log.error('The programmers are monkeys - please give them bananas ..')
						self.log.error('the mode was spelled : [%s]' % mode)
						self.log.error('.. if it works, we are lucky - but it may work.')
						mode = 'request'

					# ignore EOL
					r_buffer = r_buffer.lstrip('\r\n')

					# check to see if we have read an entire request
					request, r_buffer, seek = self.checkRequest(r_buffer, seek)

					if not request:
						yield '', ''
						continue
					seek = 0
					processing = True

					# nb_to_send is how much we expect to need to get the rest of the request
					mode, nb_to_send = yield request, ''

				# break out of the outer loop as soon as we leave the inner loop
				# through normal execution
				break

			except socket.error, e:
				if e.args[0] in errno_block:
					yield '', ''
				else:
					break

		yield None,None

	def setPeer (self, peer):
		"""Set the claimed ip address for this client.
		Does not effect the ip address we try sending data to."""
		self.peer = peer

	def readData(self):
		request,content = self.reader.send(('transfer',0))
		return self.name, self.peer, request, content

	def readRelated(self, mode, remaining):
		request, content = self.reader.send((mode,remaining))
		return self.name, self.peer, request, content

	def _write(self, sock):
		"""Coroutine managing data sent to the client"""

		w_buffer = ''
		filename = yield None

		# check to see if we are returning data directly from a local file
		if filename is not None:
			try:
				# XXX: we must read from the file on demand rather than doing this
				with open(filename) as fd:
					w_buffer = fd.read()

				found = True, False, 0
			except IOError:
				found = None

			data = yield found
			w_buffer = data + w_buffer
		else:
			found = None

		data = yield found
		finished = False

		while True:
			try:
				while True:
					had_buffer = bool(w_buffer)

					if data is not None:
						w_buffer += data
					else:
						# We've finished downloading, even if the client hasn't yet
						finished = True

					if finished:
						if not w_buffer:
							break      # terminate the client connection
						elif data:
							self.log.error('Tried to send data to client after we told it to close. Dropping it.')

					if not had_buffer or data == '':
						sent = sock.send(w_buffer)
						if sent:
							self.log.debug(">> [%s]" % w_buffer[:sent].replace('\t','\\t').replace('\r','\\r').replace('\n','\\n'))
						w_buffer = w_buffer[sent:]
					else:
						sent = 0

					buffered = bool(w_buffer) or finished
					data = yield buffered, had_buffer, sent

				# break out of the outer loop as soon as we leave the inner loop
				# through normal execution
				yield None
				break

			except socket.error, e:
				if e.args[0] in errno_block:
					#self.log.info('interrupted when trying to sent %d bytes, will retry' % len(data))
					#self.log.info('reason: errno %d: %s' % (e.args[0], errno.errorcode.get(e.args[0], '<no errno name>')))
					data = yield bool(w_buffer) or finished, had_buffer, 0
				else:
					#self.log.critical('unexpected error writing on socket')
					#self.log.critical('reason, errno %d: %s' % (e.args[0], errno.errorcode.get(e.args[0], '<no errno name>')))
					yield None  # stop the client connection
					break  # and don't come back

		yield None

	def writeData(self, data):
		res = self.writer.send(data)
		return res

	def startData(self, command, data):
		# start the _write coroutine
		self.writer.next()

		if command == 'stream':
			self.writer.send(None)  # no local file
			res = self.writer.send(data)

		elif command == 'close':
			self.writer.send(None)  # no local file
			self.writer.send(data)
			res = self.writer.send(None)  # close the connection once the buffer is empty

		elif command == 'file':
			header, filename = data
			res = self.writer.send(filename)  # use local file
			self.writer.send(header)  # write the response headers before the file

			self.writer.send(None)  # close the connection once the buffer is empty
		else:
			res = None

		# buffered, had_buffer
		return res

	def restartData(self, command, data):
		self.writer = self._write(self.sock)
		return self.startData(command, data)

	def shutdown(self):
		try:
			self.sock.shutdown(socket.SHUT_RDWR)
			self.sock.close()
		except socket.error:
			pass

		self.writer.close()
		self.reader.close()
<|endoftext|>


<commit_before>import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

class facility(object):
    """ A Facility currently under monitoring

    """

    def __init__(self , data) :
        self.validated_data = data

    def monitor_new_report(self) :
        out = np.random.choice(['Validate' , 'Supervise - Data' , 'Supervise - Services' , 'Supervise - Data and Quality'])
        return out

u = facility('a')

u.monitor_new_report()
<commit_msg>Add report making for facility<commit_after>import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from reports_monitoring import *

store = pd.HDFStore('../../data/processed/orbf_benin.h5')
data_orbf = store['data']
store.close()

class facility(object):
    """ A Facility currently under monitoring

    """

    def __init__(self , data) :
        self.validated_data = data

    def monitor_new_report(self) :
        out = np.random.choice(['Validate' , 'Supervise - Data' , 'Supervise - Services' , 'Supervise - Data and Quality'])
        return out

    def make_reports(self):
        reports = {}
        for month in list(data.date.unique()) :
            reports[str(month)[:7]] = report(data[data.date == month])
        self.reports = reports
        return reports
<|endoftext|>