<commit_before>"""
Copyright (c) 2013, Dave Mankoff
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Dave Mankoff nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL DAVE MANKOFF BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

from __future__ import unicode_literals
import sys
try:
  from html import escape
except ImportError:
  from cgi import escape

from io import StringIO

import re
try:
  from html.parser import HTMLParser
except ImportError:
  from HTMLParser import HTMLParser

PRE_TAGS = ('pre', 'textarea')  # styles and scripts are never minified
# http://www.w3.org/TR/html51/syntax.html#elements-0
NO_CLOSE_TAGS = ('area', 'base', 'br', 'col', 'command', 'embed', 'hr', 'img',
                 'input', 'keygen', 'link', 'meta', 'param', 'source', 'track',
                 'wbr')
# http://www.w3.org/TR/html51/index.html#attributes-1
BOOLEAN_ATTRIBUTES = {
  'audio': ('autoplay', 'controls', 'loop', 'muted',),
  'button': ('autofocus', 'disabled', 'formnovalidate',),
  'command': ('checked', 'disabled',),
  'dialog': ('open',),
  'fieldset': ('disabled',),
  'form': ('novalidate',),
  'iframe': ('seamless',),
  'img': ('ismap',),
  'input': ('autofocus', 'checked', 'disabled', 'formnovalidate', 'multiple',
            'readonly', 'required',),
  'keygen': ('autofocus', 'disabled',),
  'object': ('typesmustmatch',),
  'ol': ('reversed',),
  'optgroup': ('disabled',),
  'option': ('disabled', 'selected',),
  'script': ('async', 'defer'),
  'select': ('autofocus', 'disabled', 'multiple', 'required',),
  'style': ('scoped',),
  'textarea': ('autofocus', 'disabled', 'readonly', 'required',),
  'track': ('default',),
  'video': ('autoplay', 'controls', 'loop', 'muted',),
  '*': ('hidden',),
}

leading_whitespace_re = re.compile(r'^\s+')
trailing_whitespace_re = re.compile(r'\s+$')
leading_trailing_whitespace_re = re.compile(r'(^\s+)|(\s+$)')
whitespace_re = re.compile(r'\s+')
whitespace_newline_re = re.compile(r'\s*(\r|\n)+\s*')

# Tag omission rules:
# http://www.w3.org/TR/html51/syntax.html#optional-tags

class HTMLMinError(Exception): pass
class ParseError(HTMLMinError): pass
class OpenTagNotFoundError(ParseError): pass

class HTMLMinParser(HTMLParser):
  def __init__(self,
               remove_comments=False,
               remove_empty_space=False,
               remove_all_empty_space=False,
               reduce_empty_attributes=True,
               reduce_boolean_attributes=False,
               remove_optional_attribute_quotes=True,
               keep_pre=False,
               pre_tags=PRE_TAGS,
               pre_attr='pre'):
    HTMLParser.__init__(self)
    self.keep_pre = keep_pre
    self.pre_tags = pre_tags
    self.remove_comments = remove_comments
    self.remove_empty_space = remove_empty_space
    self.remove_all_empty_space = remove_all_empty_space
    self.reduce_empty_attributes = reduce_empty_attributes
    self.reduce_boolean_attributes = reduce_boolean_attributes
    self.remove_optional_attribute_quotes = remove_optional_attribute_quotes
    self.pre_attr = pre_attr
    self._data_buffer = []
    self._in_pre_tag = 0
    self._in_head = False
    self._in_title = False
    self._after_doctype = False
    self._tag_stack = []
    self._title_newly_opened = False
    self.__title_trailing_whitespace = False

  def _has_pre(self, attrs):
    for k,v in attrs:
      if k == self.pre_attr:
        return True
    return False

  def build_tag(self, tag, attrs, close_tag):
    result = StringIO()
    result.write('<')
    result.write(escape(tag))
    needs_closing_space = False
    for k,v in attrs:
      result.write(' ')
      result.write(escape(k))
      if v:
        if self.reduce_boolean_attributes and (
             k in BOOLEAN_ATTRIBUTES.get(tag,[]) or
             k in BOOLEAN_ATTRIBUTES['*']):
          pass
        elif self.remove_optional_attribute_quotes and not any((c in v for c in ('"', "'", ' ', '<', '>', '='))):
          result.write('=')
          result.write(escape(v, quote=True))
          needs_closing_space = v.endswith('/')
        else:
          result.write('="')
          result.write(escape(v, quote=True).replace(''', "'"))
          result.write('"')
      elif not self.reduce_empty_attributes:
        result.write('=""')
    if needs_closing_space:
      result.write(' ')
    if close_tag:
      result.write('/>')
    else:
      result.write('>')
    return result.getvalue()

  def handle_decl(self, decl):
    if (len(self._data_buffer) == 1 and
        whitespace_re.match(self._data_buffer[0])):
      self._data_buffer = []
    self._data_buffer.append('<!' + decl + '>\n')
    self._after_doctype = True

  def in_tag(self, *tags):
    for t in self._tag_stack:
      if t[0] in tags:
        return t
    return False

  def _close_tags_up_to(self, tag):
    num_pres = 0
    i = 0
    for i, t in enumerate(self._tag_stack):
      if t[1]:
        num_pres += 1
      if t[0] == tag:
        break

      # Only the html tag can close out everything. Put on the brakes if
      # we encounter a closing tag that we didn't recognize.
      if tag != 'html' and t[0] in ('body', 'html', 'head'):
        raise OpenTagNotFoundError()

    self._tag_stack = self._tag_stack[i+1:]

    return num_pres

  def handle_starttag(self, tag, attrs):
    self._after_doctype = False
    if tag == 'head':
      self._in_head = True
    elif self._in_head and tag == 'title':
      self._in_title = True
      self._title_newly_opened = True

    tag_sets = ( # a list of tags and tags that they are closed by
      (('li',), ('li',)),
      (('dd', 'dt',), ('dd', 'dt',)),
      (('rp', 'rt',), ('rp', 'rt',)),
      (('p',), ('address', 'article', 'aside', 'blockquote', 'dir', 'div',
                'dl', 'fieldset', 'footer', 'form', 'h1', 'h2', 'h3', 'h4',
                'h5', 'h6', 'header', 'hgroup', 'hr', 'menu', 'nav', 'ol', 'p',
                'pre', 'section', 'table', 'ul')),
      (('optgroup',), ('optgroup',)),
      (('option',), ('option', 'optgroup')),
      (('colgroup',), ('*',)),
      (('tbody', 'thead',), ('tbody', 'tfoot')),
      (('tfoot',), ('tbody',)),
      (('tr',), ('tr',)),
      (('td', 'th'), ('td', 'th')),
      )
    for open_tags, closed_by_tags in tag_sets:
      in_tag = self.in_tag(*open_tags)
      if in_tag and (tag in closed_by_tags or '*' in closed_by_tags):
        self._in_pre_tag -= self._close_tags_up_to(in_tag[0])

    start_pre = False
    if (tag in self.pre_tags or
        tag in ('script', 'style') or
        self._has_pre(attrs) or
        self._in_pre_tag > 0):
      self._in_pre_tag += 1
      start_pre = True

    self._tag_stack.insert(0, (tag, start_pre))

    if not self.keep_pre:
      attrs = [(k,v) for k,v in attrs if k != self.pre_attr]

    self._data_buffer.append(self.build_tag(tag, attrs, False))

  def handle_endtag(self, tag):
    # According to the spec,
 tags don't get closed when a parent a
    # tag closes them. Here's some logic that addresses this.
    if tag == 'a':
      contains_p = False
      for i, t in enumerate(self._tag_stack):
        if t[0] == 'p':
          contains_p = True
        elif t[0] == 'a':
          break
      if contains_p: # the p tag, and all its children should be left open
        a_tag = self._tag_stack.pop(i)
        if a_tag[1]:
          self._in_pre_tag -= 1
    else:
      if tag == 'head':
        # TODO: Did we know that we were in an head tag?! If not, we need to
        # reminify everything to remove extra spaces.
        self._in_head = False
      elif tag == 'title':
        self._in_title = False
        self._title_newly_opened = False
      try:
        self._in_pre_tag -= self._close_tags_up_to(tag)
      except OpenTagNotFoundError:
        # Some tags don't require a start tag. Most do. Either way, we leave
        # closing tags along since they affect output. For instance, a '
'
        # results in a '

' in Chrome.
        pass
    if tag not in NO_CLOSE_TAGS:
      self._data_buffer.extend(['</', escape(tag), '>'])

  def handle_startendtag(self, tag, attrs):
    self._after_doctype = False
    if not self.keep_pre:
      attrs = [(k,v) for k,v in attrs if k != 'pre']
    self._data_buffer.append(self.build_tag(tag, attrs, tag not in NO_CLOSE_TAGS))

  def handle_comment(self, data):
    if not self.remove_comments or data[0] == '!':
      self._data_buffer.append(''.format(
          data[1:] if data[0] == '!' else data))

  def handle_data(self, data):
    if self._in_pre_tag > 0:
      self._data_buffer.append(data)
    else:
      # remove_all_empty_space matches everything. remove_empty_space only
      # matches if there's a newline involved.
      if self.remove_all_empty_space or self._in_head or self._after_doctype:
        match = whitespace_re.match(data)
        if match and match.end(0) == len(data):
          return
      elif self.remove_empty_space:
        match = whitespace_newline_re.match(data)
        if match and match.end(0) == len(data):
          return

      # if we're in the title, remove leading and trailing whitespace.
      # note that the title may be parsed in chunks if entityref's or charrefs
      # are encountered.
      if self._in_title:
        if self.__title_trailing_whitespace:
          self._data_buffer.append(' ')
        self.__title_trailing_whitespace = data[-1].isspace()
        if self._title_newly_opened:
          self._title_newly_opened = False
          data = leading_trailing_whitespace_re.sub('', data)
        else:
          data = trailing_whitespace_re.sub(
            '', leading_whitespace_re.sub(' ', data))

      data = whitespace_re.sub(' ', data)
      if not data:
        return

      if self._in_pre_tag == 0 and self._data_buffer:
        # If we're not in a pre block, its possible that we append two spaces
        # together, which we want to avoid. For instance, if we remove a comment
        # from between two blocks of text: a  c => a  c.
        if data[0] == ' ' and self._data_buffer[-1][-1] == ' ':
          data = data[1:]
          if not data:
            return
      self._data_buffer.append(data)

  def handle_entityref(self, data):
    if self._in_title:
      if not self._title_newly_opened and self.__title_trailing_whitespace:
        self._data_buffer.append(' ')
        self.__title_trailing_whitespace = False
      self._title_newly_opened = False
    self._data_buffer.append('&{};'.format(data))

  def handle_charref(self, data):
    if self._in_title:
      if not self._title_newly_opened and self.__title_trailing_whitespace:
        self._data_buffer.append(' ')
        self.__title_trailing_whitespace = False
      self._title_newly_opened = False
    self._data_buffer.append('&#{};'.format(data))

  def handle_pi(self, data):
    self._data_buffer.append('<?' + data + '>')

  def unknown_decl(self, data):
    self._data_buffer.append('<![' + data + ']>')

  def reset(self):
    self._data_buffer = []
    HTMLParser.reset(self)

  @property
  def result(self):
    return ''.join(self._data_buffer)
<commit_msg>Don't remove conditional comments<commit_after>"""
Copyright (c) 2013, Dave Mankoff
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Dave Mankoff nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL DAVE MANKOFF BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

from __future__ import unicode_literals
import sys
try:
  from html import escape
except ImportError:
  from cgi import escape

from io import StringIO

import re
try:
  from html.parser import HTMLParser
except ImportError:
  from HTMLParser import HTMLParser

PRE_TAGS = ('pre', 'textarea')  # styles and scripts are never minified
# http://www.w3.org/TR/html51/syntax.html#elements-0
NO_CLOSE_TAGS = ('area', 'base', 'br', 'col', 'command', 'embed', 'hr', 'img',
                 'input', 'keygen', 'link', 'meta', 'param', 'source', 'track',
                 'wbr')
# http://www.w3.org/TR/html51/index.html#attributes-1
BOOLEAN_ATTRIBUTES = {
  'audio': ('autoplay', 'controls', 'loop', 'muted',),
  'button': ('autofocus', 'disabled', 'formnovalidate',),
  'command': ('checked', 'disabled',),
  'dialog': ('open',),
  'fieldset': ('disabled',),
  'form': ('novalidate',),
  'iframe': ('seamless',),
  'img': ('ismap',),
  'input': ('autofocus', 'checked', 'disabled', 'formnovalidate', 'multiple',
            'readonly', 'required',),
  'keygen': ('autofocus', 'disabled',),
  'object': ('typesmustmatch',),
  'ol': ('reversed',),
  'optgroup': ('disabled',),
  'option': ('disabled', 'selected',),
  'script': ('async', 'defer'),
  'select': ('autofocus', 'disabled', 'multiple', 'required',),
  'style': ('scoped',),
  'textarea': ('autofocus', 'disabled', 'readonly', 'required',),
  'track': ('default',),
  'video': ('autoplay', 'controls', 'loop', 'muted',),
  '*': ('hidden',),
}

leading_whitespace_re = re.compile(r'^\s+')
trailing_whitespace_re = re.compile(r'\s+$')
leading_trailing_whitespace_re = re.compile(r'(^\s+)|(\s+$)')
whitespace_re = re.compile(r'\s+')
whitespace_newline_re = re.compile(r'\s*(\r|\n)+\s*')

# Tag omission rules:
# http://www.w3.org/TR/html51/syntax.html#optional-tags

class HTMLMinError(Exception): pass
class ParseError(HTMLMinError): pass
class OpenTagNotFoundError(ParseError): pass

class HTMLMinParser(HTMLParser):
  def __init__(self,
               remove_comments=False,
               remove_empty_space=False,
               remove_all_empty_space=False,
               reduce_empty_attributes=True,
               reduce_boolean_attributes=False,
               remove_optional_attribute_quotes=True,
               keep_pre=False,
               pre_tags=PRE_TAGS,
               pre_attr='pre'):
    HTMLParser.__init__(self)
    self.keep_pre = keep_pre
    self.pre_tags = pre_tags
    self.remove_comments = remove_comments
    self.remove_empty_space = remove_empty_space
    self.remove_all_empty_space = remove_all_empty_space
    self.reduce_empty_attributes = reduce_empty_attributes
    self.reduce_boolean_attributes = reduce_boolean_attributes
    self.remove_optional_attribute_quotes = remove_optional_attribute_quotes
    self.pre_attr = pre_attr
    self._data_buffer = []
    self._in_pre_tag = 0
    self._in_head = False
    self._in_title = False
    self._after_doctype = False
    self._tag_stack = []
    self._title_newly_opened = False
    self.__title_trailing_whitespace = False

  def _has_pre(self, attrs):
    for k,v in attrs:
      if k == self.pre_attr:
        return True
    return False

  def build_tag(self, tag, attrs, close_tag):
    result = StringIO()
    result.write('<')
    result.write(escape(tag))
    needs_closing_space = False
    for k,v in attrs:
      result.write(' ')
      result.write(escape(k))
      if v:
        if self.reduce_boolean_attributes and (
             k in BOOLEAN_ATTRIBUTES.get(tag,[]) or
             k in BOOLEAN_ATTRIBUTES['*']):
          pass
        elif self.remove_optional_attribute_quotes and not any((c in v for c in ('"', "'", ' ', '<', '>', '='))):
          result.write('=')
          result.write(escape(v, quote=True))
          needs_closing_space = v.endswith('/')
        else:
          result.write('="')
          result.write(escape(v, quote=True).replace(''', "'"))
          result.write('"')
      elif not self.reduce_empty_attributes:
        result.write('=""')
    if needs_closing_space:
      result.write(' ')
    if close_tag:
      result.write('/>')
    else:
      result.write('>')
    return result.getvalue()

  def handle_decl(self, decl):
    if (len(self._data_buffer) == 1 and
        whitespace_re.match(self._data_buffer[0])):
      self._data_buffer = []
    self._data_buffer.append('<!' + decl + '>\n')
    self._after_doctype = True

  def in_tag(self, *tags):
    for t in self._tag_stack:
      if t[0] in tags:
        return t
    return False

  def _close_tags_up_to(self, tag):
    num_pres = 0
    i = 0
    for i, t in enumerate(self._tag_stack):
      if t[1]:
        num_pres += 1
      if t[0] == tag:
        break

      # Only the html tag can close out everything. Put on the brakes if
      # we encounter a closing tag that we didn't recognize.
      if tag != 'html' and t[0] in ('body', 'html', 'head'):
        raise OpenTagNotFoundError()

    self._tag_stack = self._tag_stack[i+1:]

    return num_pres

  def handle_starttag(self, tag, attrs):
    self._after_doctype = False
    if tag == 'head':
      self._in_head = True
    elif self._in_head and tag == 'title':
      self._in_title = True
      self._title_newly_opened = True

    tag_sets = ( # a list of tags and tags that they are closed by
      (('li',), ('li',)),
      (('dd', 'dt',), ('dd', 'dt',)),
      (('rp', 'rt',), ('rp', 'rt',)),
      (('p',), ('address', 'article', 'aside', 'blockquote', 'dir', 'div',
                'dl', 'fieldset', 'footer', 'form', 'h1', 'h2', 'h3', 'h4',
                'h5', 'h6', 'header', 'hgroup', 'hr', 'menu', 'nav', 'ol', 'p',
                'pre', 'section', 'table', 'ul')),
      (('optgroup',), ('optgroup',)),
      (('option',), ('option', 'optgroup')),
      (('colgroup',), ('*',)),
      (('tbody', 'thead',), ('tbody', 'tfoot')),
      (('tfoot',), ('tbody',)),
      (('tr',), ('tr',)),
      (('td', 'th'), ('td', 'th')),
      )
    for open_tags, closed_by_tags in tag_sets:
      in_tag = self.in_tag(*open_tags)
      if in_tag and (tag in closed_by_tags or '*' in closed_by_tags):
        self._in_pre_tag -= self._close_tags_up_to(in_tag[0])

    start_pre = False
    if (tag in self.pre_tags or
        tag in ('script', 'style') or
        self._has_pre(attrs) or
        self._in_pre_tag > 0):
      self._in_pre_tag += 1
      start_pre = True

    self._tag_stack.insert(0, (tag, start_pre))

    if not self.keep_pre:
      attrs = [(k,v) for k,v in attrs if k != self.pre_attr]

    self._data_buffer.append(self.build_tag(tag, attrs, False))

  def handle_endtag(self, tag):
    # According to the spec,
 tags don't get closed when a parent a
    # tag closes them. Here's some logic that addresses this.
    if tag == 'a':
      contains_p = False
      for i, t in enumerate(self._tag_stack):
        if t[0] == 'p':
          contains_p = True
        elif t[0] == 'a':
          break
      if contains_p: # the p tag, and all its children should be left open
        a_tag = self._tag_stack.pop(i)
        if a_tag[1]:
          self._in_pre_tag -= 1
    else:
      if tag == 'head':
        # TODO: Did we know that we were in an head tag?! If not, we need to
        # reminify everything to remove extra spaces.
        self._in_head = False
      elif tag == 'title':
        self._in_title = False
        self._title_newly_opened = False
      try:
        self._in_pre_tag -= self._close_tags_up_to(tag)
      except OpenTagNotFoundError:
        # Some tags don't require a start tag. Most do. Either way, we leave
        # closing tags along since they affect output. For instance, a '
'
        # results in a '

' in Chrome.
        pass
    if tag not in NO_CLOSE_TAGS:
      self._data_buffer.extend(['</', escape(tag), '>'])

  def handle_startendtag(self, tag, attrs):
    self._after_doctype = False
    if not self.keep_pre:
      attrs = [(k,v) for k,v in attrs if k != 'pre']
    self._data_buffer.append(self.build_tag(tag, attrs, tag not in NO_CLOSE_TAGS))

  def handle_comment(self, data):
    if not self.remove_comments or data[0] == '!' or re.match(r'\[if\s', data):
      self._data_buffer.append(''.format(
          data[1:] if data[0] == '!' else data))

  def handle_data(self, data):
    if self._in_pre_tag > 0:
      self._data_buffer.append(data)
    else:
      # remove_all_empty_space matches everything. remove_empty_space only
      # matches if there's a newline involved.
      if self.remove_all_empty_space or self._in_head or self._after_doctype:
        match = whitespace_re.match(data)
        if match and match.end(0) == len(data):
          return
      elif self.remove_empty_space:
        match = whitespace_newline_re.match(data)
        if match and match.end(0) == len(data):
          return

      # if we're in the title, remove leading and trailing whitespace.
      # note that the title may be parsed in chunks if entityref's or charrefs
      # are encountered.
      if self._in_title:
        if self.__title_trailing_whitespace:
          self._data_buffer.append(' ')
        self.__title_trailing_whitespace = data[-1].isspace()
        if self._title_newly_opened:
          self._title_newly_opened = False
          data = leading_trailing_whitespace_re.sub('', data)
        else:
          data = trailing_whitespace_re.sub(
            '', leading_whitespace_re.sub(' ', data))

      data = whitespace_re.sub(' ', data)
      if not data:
        return

      if self._in_pre_tag == 0 and self._data_buffer:
        # If we're not in a pre block, its possible that we append two spaces
        # together, which we want to avoid. For instance, if we remove a comment
        # from between two blocks of text: a  c => a  c.
        if data[0] == ' ' and self._data_buffer[-1][-1] == ' ':
          data = data[1:]
          if not data:
            return
      self._data_buffer.append(data)

  def handle_entityref(self, data):
    if self._in_title:
      if not self._title_newly_opened and self.__title_trailing_whitespace:
        self._data_buffer.append(' ')
        self.__title_trailing_whitespace = False
      self._title_newly_opened = False
    self._data_buffer.append('&{};'.format(data))

  def handle_charref(self, data):
    if self._in_title:
      if not self._title_newly_opened and self.__title_trailing_whitespace:
        self._data_buffer.append(' ')
        self.__title_trailing_whitespace = False
      self._title_newly_opened = False
    self._data_buffer.append('&#{};'.format(data))

  def handle_pi(self, data):
    self._data_buffer.append('<?' + data + '>')

  def unknown_decl(self, data):
    self._data_buffer.append('<![' + data + ']>')

  def reset(self):
    self._data_buffer = []
    HTMLParser.reset(self)

  @property
  def result(self):
    return ''.join(self._data_buffer)
<|endoftext|>


<commit_before>import os
import sys
import importlib
from contextlib import contextmanager

@contextmanager
def cwd_in_path():
    cwd = os.getcwd()
    if cwd in sys.path:
        yield
    else:
        sys.path.insert(0, cwd)
        try:
            yield cwd
        finally:
            try:
                sys.path.remove(cwd)
            except ValueError:  # pragma: no cover
                pass

def import_from_cwd(module):
    with cwd_in_path():
        return importlib.import_module(module)

def import_task_module(module):
    return import_from_cwd(module)

def import_task(fully_qualified_function_name):
    module_name, func_name = split_function_name(fully_qualified_function_name)
    module = import_task_module(module_name)
    return getattr(module, func_name)

def split_function_name(name):
    func_name, module_name = map(reverse_str, reverse_str(name).split('.', 1))
    return module_name, func_name

def reverse_str(s):
    return s[::-1]
<commit_msg>reorder functions<commit_after>import os
import sys
import importlib
from contextlib import contextmanager

def import_task(fully_qualified_function_name):
    module_name, func_name = split_function_name(fully_qualified_function_name)
    module = import_task_module(module_name)
    return getattr(module, func_name)

def split_function_name(name):
    reverse_str = lambda s: s[::-1]
    func_name, module_name = map(reverse_str, reverse_str(name).split('.', 1))
    return module_name, func_name

def import_task_module(module):
    return import_from_cwd(module)

def import_from_cwd(module):
    with cwd_in_path():
        return importlib.import_module(module)

@contextmanager
def cwd_in_path():
    cwd = os.getcwd()
    if cwd in sys.path:
        yield
    else:
        sys.path.insert(0, cwd)
        try:
            yield
        finally:
            try:
                sys.path.remove(cwd)
            except ValueError:
                pass
<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © %YEAR%  <>
#
# Distributed under terms of the %LICENSE% license.

"""

"""

import logging

logger = logging.getLogger(__name__)
import argparse
# conda install -c conda-forge begins
# import begin

import PyQt4
from PyQt4.QtGui import QGridLayout, QLabel,\
    QPushButton, QLineEdit, QApplication, QWidget, QGridLayout, QSpinBox, QLineEdit, QCheckBox,\
        QComboBox, QTextEdit, QDialog, QMainWindow, QDoubleSpinBox, QTabWidget

from PyQt4 import QtGui
import sys
import os.path as op
import copy
import inspect
import collections
import numpy as np
import scipy

import dictwidgetqt
import iowidgetqt
import dictwidgetpyqtgraph
import generators.cylinders
import generators.gensei_wrapper

from pyqtconfig import ConfigManager

class TeigenWidget(QtGui.QWidget):
    def __init__(self, ncols=2):
        super(TeigenWidget, self).__init__()
        self.ncols = ncols
        self.gen = None
        self.dataframes = {}
        self.figures = {}
        self.ui_stats_shown = False
        self.teigen = Teigen()
        self.version = "0.1.20"
        self.init_ui()

    def run(self):
        print "generator args"
        id = self.gen_tab_wg.currentIndex()
        new_cfg = self._ui_generator_widgets[id].config_as_dict()
        logger.debug(str(new_cfg))

        none, area_cfg = dictwidgetpyqtgraph.from_pyqtgraph_struct(self.area_sampling_params.saveState())

        new_cfg.update(area_cfg["Area Sampling"])

        # self.config = new_cfg
        self.teigen.run(generator_id=id, **self.config)

    def _area_sampling_gensei_export(self, area_sampling_params):
        asp = area_sampling_params
        vs_mm = np.asarray(asp["voxelsize_mm"])
        resolution = 1.0 / vs_mm
        dct = {
            'dims': asp["areasize_mm"],
            'n_slice': asp["areasize_px"][0],
            'resolution': [resolution[1], resolution[2]]
        }
        return dct

    def _area_sampling_cylinder_generator_export(self, area_sampling_params):
        return {
            'voxelsize_mm': area_sampling_params["voxelsize_mm"],
            'area_shape': area_sampling_params["areasize_px"]
        }

    def _show_stats(self):
        if self.ui_stats_shown:
        #     self._wg_tab_describe.deleteLater()
        #     self._wg_tab_describe = None
        #     self._wg_tab_merne.deleteLater()
        #     self._wg_tab_merne = None
            pass
        else:

            self.stats_tab_wg = QTabWidget()
            self.mainLayout.addWidget(self.stats_tab_wg, 0, 3, 5, 2)
        if True:
            from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
            # from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
            import matplotlib.pyplot as plt

            self.figure = plt.figure()
            self.canvas = FigureCanvas(self.figure)
            # self.toolbar = NavigationToolbar(self.canvas, self)
            self.stats_tab_wg.addTab(self.canvas, 'Graphs')

        df = self.teigen.gen.getStats()
        import tablewidget
        to_rename = {
            "length": "length [mm]",
            "volume": "volume [mm^3]",
            "surface": "surface [mm^2]",
            "radius": "radius [mm^2]"
        }
        to_rename_density = {
            "length": "length [mm^-2]",
            "volume": "volume []",
            "surface": "surface [mm^-1]",
            "radius": "radius [mm^-2]"
        }

        dfmerne = df[["length", "volume", "surface", "radius"]].sum() / self.gen.area_volume
        print "merne"
        print dfmerne
        dfmernef = dfmerne.to_frame().transpose().rename(columns=to_rename_density)
        # dfmernef.insert(0, "", dfmernef.index)
        # import ipdb; ipdb.set_trace()

        self._wg_tab_merne = tablewidget.TableWidget(self, dataframe=dfmernef)
        self.stats_tab_wg.addTab(self._wg_tab_merne, "Density table")
        self.dataframes["density"] = dfmernef

        plt.subplot(141)
        df[["length"]].rename(columns=to_rename).boxplot(return_type='axes')
        plt.subplot(142)
        df[['radius']].rename(columns=to_rename).boxplot(return_type='axes')

        plt.subplot(143)
        df[["surface"]].rename(columns=to_rename).boxplot(return_type='axes')

        plt.subplot(144)
        df[["volume"]].rename(columns=to_rename).boxplot(return_type='axes')

        # TODO take care about redrawing
        dfdescribe = df.describe()
        dfdescribe.insert(0, "", dfdescribe.index)
        self.dataframes["describe"] = dfdescribe

        self._wg_tab_describe = tablewidget.TableWidget(self, dataframe=dfdescribe)
        self._wg_tab_describe.show()
        self._wg_tab_describe.raise_()
        self._wg_tab_describe.setMinimumWidth(600)
        self._wg_tab_describe.setMinimumHeight(200)

        # self.mainLayout.addWidget(self._wg_tab_describe, 0, 2, 5, 2)
        self.stats_tab_wg.addTab(self._wg_tab_describe, "Stats table")
        self.resize(600,700)

        self.ui_stats_shown = True

    def complicated_to_yaml(self, cfg):
        import yaml
        # convert values to json
        isconverted = {}
        for key, value in cfg.iteritems():
            if type(value) in (str, int, float, bool):

                isconverted[key] = False
                if type(value) is str:
                    pass

            else:
                isconverted[key] = True
                cfg[key] = yaml.dump(value, default_flow_style=True)
        return cfg

    def _get_generator(self, id):
        pass

    def init_ui(self):
        wtitle = "Teigen " + self.version
        self.setWindowTitle(wtitle)
        self.mainLayout = QGridLayout(self)

        hide_keys = ["build", "gtree"]
        self.gen_tab_wg = QTabWidget()
        self.mainLayout.addWidget(self.gen_tab_wg, 0, 1)

        rename_captions_dict = {
            "voxelsize_mm": "voxel size [mm]",
            }

        self._ui_generator_widgets = []
        for i, config in enumerate(self.teigen.configs):
            wg = dictwidgetqt.DictWidget(
                self.teigen.configs[i],
                hide_keys=hide_keys,
                captions=rename_captions_dict,
                ncols=1,
            )
            self._ui_generator_widgets.append(wg)
            self.gen_tab_wg.addTab(wg, self.generators_names[i])
        # self.gen_tab_wg.addTab(gen_wg, "cylinder generator")
        # self.gen_tab_wg.addTab(gen_wg, "gensei generator")

        # self.mainLayout.setColumnMinimumWidth(text_col, 500)

        postprocessing_params = dictwidgetqt.get_default_args(self.teigen.postprocessing)
        self.posprocessing_wg = dictwidgetqt.DictWidget(postprocessing_params)
        self.mainLayout.addWidget(self.posprocessing_wg, 1, 1)

        btn_accept = QPushButton("Run", self)
        btn_accept.clicked.connect(self.btnRun)
        self.mainLayout.addWidget(btn_accept, 2, 1) # , (gd_max_i / 2), text_col)

        self.ui_output_dir_widget = iowidgetqt.SetDirWidget("~/teigen_data/slice{:06d}.jpg", "output directory")
        self.ui_output_dir_widget.setToolTip("Data are stored in defined directory.\nOutput format is based on file extension.\nFor saving into image stack use 'filename{:06d}.jpg'")
        self.mainLayout.addWidget(self.ui_output_dir_widget, 3, 1) # , (gd_max_i / 2), text_col)

        btn_save = QPushButton("Save", self)
        btn_save.setToolTip("Save image slices and meta information")
        btn_save.clicked.connect(self.btnSave)
        self.mainLayout.addWidget(btn_save, 3, 1) # , (gd_max_i / 2), text_col)
        # self.config.updated.connect(self.on_config_update)

        ## pyqtgraph experiments
        import dictwidgetpyqtgraph
        import pyqtgraph
        from pyqtgraph.parametertree import Parameter, ParameterTree, ParameterItem, registerParameterType
        input_params = {
            "Area Sampling":  dictwidgetpyqtgraph.AreaSamplingParameter(name='Area Sampling')
        }
        gr_struct = dictwidgetpyqtgraph.to_pyqtgraph_struct('params', input_params, opts={})
        p = Parameter.create(**gr_struct)

        t = ParameterTree()
        t.setParameters(p, showTop=False)
        t.setMinimumWidth(300)
        t.show()
        self.mainLayout.addWidget(t, 0, 0, 5, 1)
        self.area_sampling_wg = t
        self.area_sampling_params = p

        ## end of pyqtgraph tree
    def btnRun(self):

        logger.debug("btnAccept")
        logger.debug(str(self.config))
        self.run()
        self._show_stats()

    def btnSave(self):
        # filename = "file{:05d}.jpg"
        # filename = "file%05d.jpg"
        # filename = QtGui.QFileDialog.getSaveFileName(
        #     self,
        #     "Save file",
        #     init_filename,
        #     ""
        # )
        # filename = str(filename)

        if self.gen is None:
            self.run()
            self._show_stats()

        # postprocessing
        if "generate_volume" in dir(self.gen):
            self.teigen.data3d = self.gen.generate_volume()
            self.teigen.voxelsize_mm = self.gen.voxelsize_mm
            postprocessing_params = self.posprocessing_wg.config_as_dict()
            data3d = self.teigen.postprocessing(**postprocessing_params)
            self.gen.data3d = data3d

        # filename = op.join(self.ui_output_dir_widget.get_dir(), filename)
        filename = self.ui_output_dir_widget.get_dir()

        filename = iowidgetqt.str_format_old_to_new(filename)

        self.gen.saveVolumeToFile(filename)
        fn_base = self.filename_base(filename)
        for dfname in self.dataframes:
            df = self.dataframes[dfname].to_csv(fn_base+"_" + dfname + ".csv")
        self.figure.savefig(fn_base + "_" + "graph.pdf")
        self.figure.savefig(fn_base + "_" + "graph.png")
        self.figure.savefig(fn_base + "_" + "graph.svg")

    def filename_base(self, filename):

        import re
        filename = re.sub(r"({.*})", r"", filename)

        root, ext = op.splitext(filename)
        return root

    def on_config_update(self):
        pass

    def config_as_dict(self):
        dictionary = self.config.as_dict()
        for key, value in dictionary.iteritems():
            from PyQt4.QtCore import pyqtRemoveInputHook
            pyqtRemoveInputHook()
            # import ipdb; ipdb.set_trace() #  noqa BREAKPOINT
            if type(value) == PyQt4.QtCore.QString:
                value = str(value)
            dictionary[key] = value

        return dictionary

class Teigen():
    def __init__(self):
        self.data3d = None
        self.voxelsize_mm = None
        self.generators_classes =  [
            generators.cylinders.CylinderGenerator,
            generators.gensei_wrapper.GenseiGenerator
        ]
        self.generators_names = [
            "Cylinder generator",
            "Gensei generator"
        ]
        self.configs = [dictwidgetqt.get_default_args(conf) for conf in self.generators_classes]
        self.config = self.configs[0]

    def run(self, **config):
        id = config.pop('id')

        self.config = config
        generator_class = self.generators_classes[id]
        # self.config = get_default_args(generator_class)
        self.gen = generator_class(**self.config)
        # self.gen = generators.gensei_wrapper.GenseiGenerator(**self.config2)
        # self.gen = generators.gensei_wrapper.GenseiGenerator()
        self.gen.run()

    def postprocessing(self, gaussian_filter=True, gaussian_filter_sigma_mm=1.0, gaussian_noise=True, gaussian_noise_stddev=10.0, gaussian_noise_center=0.0, limit_negative_intensities=True):
        if gaussian_filter:
            sigma_px = gaussian_filter_sigma_mm / self.voxelsize_mm

            self.data3d = scipy.ndimage.filters.gaussian_filter(
                self.data3d,
                sigma=sigma_px)

        if gaussian_noise:
            dt = self.data3d.dtype
            noise = np.random.normal(loc=gaussian_noise_center, scale=gaussian_noise_stddev, size=self.data3d.shape)
            self.data3d = (self.data3d + noise).astype(self.data3d.dtype)

        if limit_negative_intensities:
            self.data3d[self.data3d < 0] = 0

        return self.data3d

    # def save_volume_to_file(self, filename):
    #
    #     import io3
    #     import io3d.misc
    #     import numpy as np
    #     data = {
    #         'data3d': self.data3d.astype(np.uint8), #* self.output_intensity,
    #         'voxelsize_mm': self.voxelsize_mm,
    #         # 'segmentation': np.zeros_like(self.data3d, dtype=np.int8)
    #     }
    #     io3d.write(data, filename)

def main():
    logger = logging.getLogger()

    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    logger.addHandler(ch)

    # create file handler which logs even debug messages
    # fh = logging.FileHandler('log.txt')
    # fh.setLevel(logging.DEBUG)
    # formatter = logging.Formatter(
    #     '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    # fh.setFormatter(formatter)
    # logger.addHandler(fh)
    # logger.debug('start')

    # input parser
    parser = argparse.ArgumentParser(
        description=__doc__
    )
    parser.add_argument(
        '-i', '--inputfile',
        default=None,
        # required=True,
        help='input file'
    )
    parser.add_argument(
        '-d', '--debug', action='store_true',
        help='Debug mode')
    args = parser.parse_args()

    if args.debug:
        ch.setLevel(logging.DEBUG)

    app = QApplication(sys.argv)
    cw = TeigenWidget()
    cw.show()
    app.exec_()

if __name__ == "__main__":
    main()<commit_msg>fixed calling of moved list<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © %YEAR%  <>
#
# Distributed under terms of the %LICENSE% license.

"""

"""

import logging

logger = logging.getLogger(__name__)
import argparse
# conda install -c conda-forge begins
# import begin

import PyQt4
from PyQt4.QtGui import QGridLayout, QLabel,\
    QPushButton, QLineEdit, QApplication, QWidget, QGridLayout, QSpinBox, QLineEdit, QCheckBox,\
        QComboBox, QTextEdit, QDialog, QMainWindow, QDoubleSpinBox, QTabWidget

from PyQt4 import QtGui
import sys
import os.path as op
import copy
import inspect
import collections
import numpy as np
import scipy

import dictwidgetqt
import iowidgetqt
import dictwidgetpyqtgraph
import generators.cylinders
import generators.gensei_wrapper

from pyqtconfig import ConfigManager

class TeigenWidget(QtGui.QWidget):
    def __init__(self, ncols=2):
        super(TeigenWidget, self).__init__()
        self.ncols = ncols
        self.gen = None
        self.dataframes = {}
        self.figures = {}
        self.ui_stats_shown = False
        self.teigen = Teigen()
        self.version = "0.1.20"
        self.init_ui()

    def run(self):
        print "generator args"
        id = self.gen_tab_wg.currentIndex()
        new_cfg = self._ui_generator_widgets[id].config_as_dict()
        logger.debug(str(new_cfg))

        none, area_cfg = dictwidgetpyqtgraph.from_pyqtgraph_struct(self.area_sampling_params.saveState())

        new_cfg.update(area_cfg["Area Sampling"])

        # self.config = new_cfg
        self.teigen.run(generator_id=id, **self.config)

    def _area_sampling_gensei_export(self, area_sampling_params):
        asp = area_sampling_params
        vs_mm = np.asarray(asp["voxelsize_mm"])
        resolution = 1.0 / vs_mm
        dct = {
            'dims': asp["areasize_mm"],
            'n_slice': asp["areasize_px"][0],
            'resolution': [resolution[1], resolution[2]]
        }
        return dct

    def _area_sampling_cylinder_generator_export(self, area_sampling_params):
        return {
            'voxelsize_mm': area_sampling_params["voxelsize_mm"],
            'area_shape': area_sampling_params["areasize_px"]
        }

    def _show_stats(self):
        if self.ui_stats_shown:
        #     self._wg_tab_describe.deleteLater()
        #     self._wg_tab_describe = None
        #     self._wg_tab_merne.deleteLater()
        #     self._wg_tab_merne = None
            pass
        else:

            self.stats_tab_wg = QTabWidget()
            self.mainLayout.addWidget(self.stats_tab_wg, 0, 3, 5, 2)
        if True:
            from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
            # from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
            import matplotlib.pyplot as plt

            self.figure = plt.figure()
            self.canvas = FigureCanvas(self.figure)
            # self.toolbar = NavigationToolbar(self.canvas, self)
            self.stats_tab_wg.addTab(self.canvas, 'Graphs')

        df = self.teigen.gen.getStats()
        import tablewidget
        to_rename = {
            "length": "length [mm]",
            "volume": "volume [mm^3]",
            "surface": "surface [mm^2]",
            "radius": "radius [mm^2]"
        }
        to_rename_density = {
            "length": "length [mm^-2]",
            "volume": "volume []",
            "surface": "surface [mm^-1]",
            "radius": "radius [mm^-2]"
        }

        dfmerne = df[["length", "volume", "surface", "radius"]].sum() / self.gen.area_volume
        print "merne"
        print dfmerne
        dfmernef = dfmerne.to_frame().transpose().rename(columns=to_rename_density)
        # dfmernef.insert(0, "", dfmernef.index)
        # import ipdb; ipdb.set_trace()

        self._wg_tab_merne = tablewidget.TableWidget(self, dataframe=dfmernef)
        self.stats_tab_wg.addTab(self._wg_tab_merne, "Density table")
        self.dataframes["density"] = dfmernef

        plt.subplot(141)
        df[["length"]].rename(columns=to_rename).boxplot(return_type='axes')
        plt.subplot(142)
        df[['radius']].rename(columns=to_rename).boxplot(return_type='axes')

        plt.subplot(143)
        df[["surface"]].rename(columns=to_rename).boxplot(return_type='axes')

        plt.subplot(144)
        df[["volume"]].rename(columns=to_rename).boxplot(return_type='axes')

        # TODO take care about redrawing
        dfdescribe = df.describe()
        dfdescribe.insert(0, "", dfdescribe.index)
        self.dataframes["describe"] = dfdescribe

        self._wg_tab_describe = tablewidget.TableWidget(self, dataframe=dfdescribe)
        self._wg_tab_describe.show()
        self._wg_tab_describe.raise_()
        self._wg_tab_describe.setMinimumWidth(600)
        self._wg_tab_describe.setMinimumHeight(200)

        # self.mainLayout.addWidget(self._wg_tab_describe, 0, 2, 5, 2)
        self.stats_tab_wg.addTab(self._wg_tab_describe, "Stats table")
        self.resize(600,700)

        self.ui_stats_shown = True

    def complicated_to_yaml(self, cfg):
        import yaml
        # convert values to json
        isconverted = {}
        for key, value in cfg.iteritems():
            if type(value) in (str, int, float, bool):

                isconverted[key] = False
                if type(value) is str:
                    pass

            else:
                isconverted[key] = True
                cfg[key] = yaml.dump(value, default_flow_style=True)
        return cfg

    def _get_generator(self, id):
        pass

    def init_ui(self):
        wtitle = "Teigen " + self.version
        self.setWindowTitle(wtitle)
        self.mainLayout = QGridLayout(self)

        hide_keys = ["build", "gtree"]
        self.gen_tab_wg = QTabWidget()
        self.mainLayout.addWidget(self.gen_tab_wg, 0, 1)

        rename_captions_dict = {
            "voxelsize_mm": "voxel size [mm]",
            }

        self._ui_generator_widgets = []
        for i, config in enumerate(self.teigen.configs):
            wg = dictwidgetqt.DictWidget(
                self.teigen.configs[i],
                hide_keys=hide_keys,
                captions=rename_captions_dict,
                ncols=1,
            )
            self._ui_generator_widgets.append(wg)
            self.gen_tab_wg.addTab(wg, self.teigen.generators_names[i])
        # self.gen_tab_wg.addTab(gen_wg, "cylinder generator")
        # self.gen_tab_wg.addTab(gen_wg, "gensei generator")

        # self.mainLayout.setColumnMinimumWidth(text_col, 500)

        postprocessing_params = dictwidgetqt.get_default_args(self.teigen.postprocessing)
        self.posprocessing_wg = dictwidgetqt.DictWidget(postprocessing_params)
        self.mainLayout.addWidget(self.posprocessing_wg, 1, 1)

        btn_accept = QPushButton("Run", self)
        btn_accept.clicked.connect(self.btnRun)
        self.mainLayout.addWidget(btn_accept, 2, 1) # , (gd_max_i / 2), text_col)

        self.ui_output_dir_widget = iowidgetqt.SetDirWidget("~/teigen_data/slice{:06d}.jpg", "output directory")
        self.ui_output_dir_widget.setToolTip("Data are stored in defined directory.\nOutput format is based on file extension.\nFor saving into image stack use 'filename{:06d}.jpg'")
        self.mainLayout.addWidget(self.ui_output_dir_widget, 3, 1) # , (gd_max_i / 2), text_col)

        btn_save = QPushButton("Save", self)
        btn_save.setToolTip("Save image slices and meta information")
        btn_save.clicked.connect(self.btnSave)
        self.mainLayout.addWidget(btn_save, 3, 1) # , (gd_max_i / 2), text_col)
        # self.config.updated.connect(self.on_config_update)

        ## pyqtgraph experiments
        import dictwidgetpyqtgraph
        import pyqtgraph
        from pyqtgraph.parametertree import Parameter, ParameterTree, ParameterItem, registerParameterType
        input_params = {
            "Area Sampling":  dictwidgetpyqtgraph.AreaSamplingParameter(name='Area Sampling')
        }
        gr_struct = dictwidgetpyqtgraph.to_pyqtgraph_struct('params', input_params, opts={})
        p = Parameter.create(**gr_struct)

        t = ParameterTree()
        t.setParameters(p, showTop=False)
        t.setMinimumWidth(300)
        t.show()
        self.mainLayout.addWidget(t, 0, 0, 5, 1)
        self.area_sampling_wg = t
        self.area_sampling_params = p

        ## end of pyqtgraph tree
    def btnRun(self):

        logger.debug("btnAccept")
        logger.debug(str(self.config))
        self.run()
        self._show_stats()

    def btnSave(self):
        # filename = "file{:05d}.jpg"
        # filename = "file%05d.jpg"
        # filename = QtGui.QFileDialog.getSaveFileName(
        #     self,
        #     "Save file",
        #     init_filename,
        #     ""
        # )
        # filename = str(filename)

        if self.gen is None:
            self.run()
            self._show_stats()

        # postprocessing
        if "generate_volume" in dir(self.gen):
            self.teigen.data3d = self.gen.generate_volume()
            self.teigen.voxelsize_mm = self.gen.voxelsize_mm
            postprocessing_params = self.posprocessing_wg.config_as_dict()
            data3d = self.teigen.postprocessing(**postprocessing_params)
            self.gen.data3d = data3d

        # filename = op.join(self.ui_output_dir_widget.get_dir(), filename)
        filename = self.ui_output_dir_widget.get_dir()

        filename = iowidgetqt.str_format_old_to_new(filename)

        self.gen.saveVolumeToFile(filename)
        fn_base = self.filename_base(filename)
        for dfname in self.dataframes:
            df = self.dataframes[dfname].to_csv(fn_base+"_" + dfname + ".csv")
        self.figure.savefig(fn_base + "_" + "graph.pdf")
        self.figure.savefig(fn_base + "_" + "graph.png")
        self.figure.savefig(fn_base + "_" + "graph.svg")

    def filename_base(self, filename):

        import re
        filename = re.sub(r"({.*})", r"", filename)

        root, ext = op.splitext(filename)
        return root

    def on_config_update(self):
        pass

    def config_as_dict(self):
        dictionary = self.config.as_dict()
        for key, value in dictionary.iteritems():
            from PyQt4.QtCore import pyqtRemoveInputHook
            pyqtRemoveInputHook()
            # import ipdb; ipdb.set_trace() #  noqa BREAKPOINT
            if type(value) == PyQt4.QtCore.QString:
                value = str(value)
            dictionary[key] = value

        return dictionary

class Teigen():
    def __init__(self):
        self.data3d = None
        self.voxelsize_mm = None
        self.generators_classes =  [
            generators.cylinders.CylinderGenerator,
            generators.gensei_wrapper.GenseiGenerator
        ]
        self.generators_names = [
            "Cylinder generator",
            "Gensei generator"
        ]
        self.configs = [dictwidgetqt.get_default_args(conf) for conf in self.generators_classes]
        self.config = self.configs[0]

    def run(self, **config):
        id = config.pop('id')

        self.config = config
        generator_class = self.generators_classes[id]
        # self.config = get_default_args(generator_class)
        self.gen = generator_class(**self.config)
        # self.gen = generators.gensei_wrapper.GenseiGenerator(**self.config2)
        # self.gen = generators.gensei_wrapper.GenseiGenerator()
        self.gen.run()

    def postprocessing(self, gaussian_filter=True, gaussian_filter_sigma_mm=1.0, gaussian_noise=True, gaussian_noise_stddev=10.0, gaussian_noise_center=0.0, limit_negative_intensities=True):
        if gaussian_filter:
            sigma_px = gaussian_filter_sigma_mm / self.voxelsize_mm

            self.data3d = scipy.ndimage.filters.gaussian_filter(
                self.data3d,
                sigma=sigma_px)

        if gaussian_noise:
            dt = self.data3d.dtype
            noise = np.random.normal(loc=gaussian_noise_center, scale=gaussian_noise_stddev, size=self.data3d.shape)
            self.data3d = (self.data3d + noise).astype(self.data3d.dtype)

        if limit_negative_intensities:
            self.data3d[self.data3d < 0] = 0

        return self.data3d

    # def save_volume_to_file(self, filename):
    #
    #     import io3
    #     import io3d.misc
    #     import numpy as np
    #     data = {
    #         'data3d': self.data3d.astype(np.uint8), #* self.output_intensity,
    #         'voxelsize_mm': self.voxelsize_mm,
    #         # 'segmentation': np.zeros_like(self.data3d, dtype=np.int8)
    #     }
    #     io3d.write(data, filename)

def main():
    logger = logging.getLogger()

    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    logger.addHandler(ch)

    # create file handler which logs even debug messages
    # fh = logging.FileHandler('log.txt')
    # fh.setLevel(logging.DEBUG)
    # formatter = logging.Formatter(
    #     '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    # fh.setFormatter(formatter)
    # logger.addHandler(fh)
    # logger.debug('start')

    # input parser
    parser = argparse.ArgumentParser(
        description=__doc__
    )
    parser.add_argument(
        '-i', '--inputfile',
        default=None,
        # required=True,
        help='input file'
    )
    parser.add_argument(
        '-d', '--debug', action='store_true',
        help='Debug mode')
    args = parser.parse_args()

    if args.debug:
        ch.setLevel(logging.DEBUG)

    app = QApplication(sys.argv)
    cw = TeigenWidget()
    cw.show()
    app.exec_()

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


<commit_before># -*- coding: utf-8 -*-
import os
from os.path import join as opj
import subprocess
import sys

class UnmetDependency(OSError):
    pass

def qdiff(left, right):
    left, right = opj(left, ''), opj(right, '')
    os.chdir(left)
    subprocess.check_call(['bzr', 'init', '-q'])
    subprocess.check_call(['bzr', 'add', '-q'])
    subprocess.check_call(['bzr', 'commit', '-qmm'])
    subprocess.check_call(['rsync', '-qavLt', '--delete', '--exclude=/.bzr', right, left])
    subprocess.check_call(['bzr', 'add', '-q'])
    subprocess.check_call(['bzr', 'qdiff', '-q'])

def check_dependencies():
    try:
        gitv = subprocess.check_output(['git', '--version']).split()[2]
    except subprocess.CalledProcessError:
        raise UnmetDependency("Could not find `git` program")

    ver = map(int, gitv.split('.')[0:3])
    if cmp(ver, [1, 7, 11]) < 0:
        raise UnmetDependency("Need at least `git` version 1.7.11+")

    try:
        subprocess.check_output(['bzr', 'version']).splitlines()[0].split()[-1]
    except subprocess.CalledProcessError:
        raise UnmetDependency("Could not find `bzr` program")

    try:
        subprocess.check_output(['bzr', 'qdiff', '--help'])
    except subprocess.CalledProcessError:
        raise UnmetDependency("Could not find `bzr qdiff` plugin")

    try:
        subprocess.check_output(['rsync', '--help'])
    except subprocess.CalledProcessError:
        raise UnmetDependency("Could not find `rsync` program")

def main():
    try:
        check_dependencies()
    except Exception, e:
        sys.exit(e.message)
    if len(sys.argv) == 3:
        left, right = sys.argv[1:3]
        if os.path.isdir(left) and os.path.isdir(right):
            qdiff(left, right)
            sys.exit(0)

    # TODO: ensure infinite subcalls is not possible
    cmd = ['git', 'difftool', '-d', '-x'] + sys.argv
    # print("Launching '%s'" % ' '.join(cmd))
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError, e:
        sys.exit("git failed with returncode %d" % e.returncode)
<commit_msg>Make it work on both python2 and python3 (Fixes #4)<commit_after># -*- coding: utf-8 -*-
import os
from os.path import join as opj
import subprocess
import sys

class UnmetDependency(OSError):
    pass

class CallError(Exception):
    def __init__(self, *args, **kw):
        returncode = kw.pop('returncode', 1)
        super(CallError, self).__init__(*args, **kw)
        self.returncode = returncode

def call(*args):
    try:
        r = subprocess.check_output(args, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError as e:
        if e.output:
            raise CallError(e.output.decode('utf8'), returncode=e.returncode)
        raise e
    return r.decode('utf8')

def qdiff(left, right):
    left, right = opj(left, ''), opj(right, '')
    os.chdir(left)
    call('bzr', 'init', '-q')
    call('bzr', 'add', '-q')
    call('bzr', 'commit', '-qmm')
    call('rsync', '-qavLt', '--delete', '--exclude=/.bzr', right, left)
    call('bzr', 'add', '-q')
    call('bzr', 'qdiff', '-q')

def check_dependencies():
    try:
        gitv = call('git', '--version').split()[2]
    except CallError:
        raise UnmetDependency("Could not find `git` program")

    ver = list(map(int, gitv.split('.')[0:3]))
    if ver < [1, 7, 11]:
        raise UnmetDependency("Need at least `git` version 1.7.11+")

    try:
        call('bzr', 'version').splitlines()[0].split()[-1]
    except CallError:
        raise UnmetDependency("Could not find `bzr` program")

    try:
        call('bzr', 'qdiff', '--help')
    except CallError:
        raise UnmetDependency("Could not find `bzr qdiff` plugin")

    try:
        call('rsync', '--help')
    except CallError:
        raise UnmetDependency("Could not find `rsync` program")

def main():
    try:
        check_dependencies()
    except Exception as e:
        sys.exit(str(e))
    if len(sys.argv) == 3:
        left, right = sys.argv[1:3]
        if os.path.isdir(left) and os.path.isdir(right):
            try:
                qdiff(left, right)
            except Exception as e:
                sys.exit(str(e))
            sys.exit(0)

    # TODO: ensure infinite subcalls is not possible
    cmd = ['git', 'difftool', '-d', '-x'] + sys.argv
    # print("Launching '%s'" % ' '.join(cmd))
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        sys.stderr.write(str(e))
        sys.exit(e.returncode)
<|endoftext|>


<commit_before>from __future__ import absolute_import, print_function
import os
import sys
import json
import time
import errno
import socket
import signal
import logging
import logging.config
import warnings
import threading
import traceback
import multiprocessing

import amqp

from kuyruk import importer, signals
from kuyruk.task import get_queue_name
from kuyruk.exceptions import Reject, Discard, ConnectionError

logger = logging.getLogger(__name__)

class Worker(object):
    """Consumes tasks from a queue and runs them.

    :param app: An instance of :class:`~kuyruk.Kuyruk`
    :param args: Command line arguments

    """
    def __init__(self, app, args):
        self.kuyruk = app

        if not args.queues:
            raise ValueError("no queue given")

        self.queues = [get_queue_name(q, local=args.local) for q in args.queues]
        self.shutdown_pending = threading.Event()
        self._pause = False
        self._started = None
        self.consuming = False
        self.current_task = None
        self.current_args = None
        self.current_kwargs = None
        self._heartbeat_exc_info = None
        if self.config.WORKER_MAX_LOAD is None:
            self.config.WORKER_MAX_LOAD = multiprocessing.cpu_count()

        self._pid = os.getpid()
        self._hostname = socket.gethostname()

        signals.worker_init.send(self.kuyruk, worker=self)

    @property
    def queue(self):
        warnings.warn("Worker.queue is deprecated. Use Worker.queues instead.")
        return self.queues[0]

    @property
    def config(self):
        return self.kuyruk.config

    def run(self):
        """Runs the worker and consumes messages from RabbitMQ.
        Returns only after `shutdown()` is called.

        """
        # Lazy import setproctitle.
        # There is bug with the latest version of Python with
        # uWSGI and setproctitle combination.
        # Watch: https://github.com/unbit/uwsgi/issues/1030
        from setproctitle import setproctitle
        setproctitle("kuyruk: worker on %s" % ','.join(self.queues))

        self._setup_logging()

        signal.signal(signal.SIGINT, self._handle_sigint)
        signal.signal(signal.SIGTERM, self._handle_sigterm)
        signal.signal(signal.SIGHUP, self._handle_sighup)
        signal.signal(signal.SIGUSR1, self._handle_sigusr1)
        signal.signal(signal.SIGUSR2, self._handle_sigusr2)

        self._started = os.times()[4]

        for f in (self._watch_load, self._shutdown_timer):
            t = threading.Thread(target=f)
            t.daemon = True
            t.start()

        signals.worker_start.send(self.kuyruk, worker=self)
        self._consume_messages()
        signals.worker_shutdown.send(self.kuyruk, worker=self)

        logger.debug("End run worker")

    def _setup_logging(self):
        if self.config.WORKER_LOGGING_CONFIG:
            logging.config.fileConfig(self.config.WORKER_LOGGING_CONFIG)
        else:
            level = getattr(logging, self.config.WORKER_LOGGING_LEVEL.upper())
            fmt = "%(levelname).1s " \
                  "%(name)s.%(funcName)s:%(lineno)d - %(message)s"
            logging.basicConfig(level=level, format=fmt)

    def _consume_messages(self):
        with self.kuyruk.channel() as ch:
            # Set prefetch count to 1. If we don't set this, RabbitMQ keeps
            # sending messages while we are already working on a message.
            ch.basic_qos(0, 1, True)

            self._declare_queues(ch)
            while not self.shutdown_pending.is_set():
                # Consume or pause
                if self._pause and self.consuming:
                    self._cancel_queues(ch)
                    logger.info('Consumer cancelled')
                    self.consuming = False
                elif not self._pause and not self.consuming:
                    self._consume_queues(ch)
                    logger.info('Consumer started')
                    self.consuming = True

                if self._heartbeat_exc_info:
                    break

                try:
                    ch.connection.heartbeat_tick()
                    ch.connection.drain_events(timeout=1)
                except socket.timeout:
                    pass
                except socket.error as e:
                    if e.errno != errno.EINTR:
                        raise
        logger.debug("End run worker")

    def _consumer_tag(self, queue):
        return "%s:%s@%s" % (queue, self._pid, self._hostname)

    def _declare_queues(self, ch):
        for queue in self.queues:
            logger.debug("queue_declare: %s", queue)
            ch.queue_declare(
                queue=queue, durable=True, auto_delete=False)

    def _consume_queues(self, ch):
        for queue in self.queues:
            logger.debug("basic_consume: %s", queue)
            ch.basic_consume(queue=queue,
                             consumer_tag=self._consumer_tag(queue),
                             callback=self._process_message)

    def _cancel_queues(self, ch):
        for queue in self.queues:
            logger.debug("basic_cancel: %s", queue)
            ch.basic_cancel(self._consumer_tag(queue))

    def _process_message(self, message):
        """Processes the message received from the queue."""
        if self.shutdown_pending.is_set():
            return

        try:
            description = json.loads(message.body)
        except Exception:
            message.channel.basic_reject(message.delivery_tag, requeue=False)
            logger.error("Cannot decode message. Dropping.")
        else:
            logger.info("Processing task: %r", description)
            self._process_description(message, description)

    def _process_description(self, message, description):
        try:
            task = importer.import_object(description['module'],
                                          description['function'])
            args, kwargs = description['args'], description['kwargs']
        except Exception:
            logger.error('Cannot import task')
            exc_info = sys.exc_info()
            signals.worker_failure.send(self.kuyruk, description=description,
                                        exc_info=exc_info, worker=self)
            message.channel.basic_reject(message.delivery_tag, requeue=False)
        else:
            self._process_task(message, description, task, args, kwargs)

    def _process_task(self, message, description, task, args, kwargs):
        reply_to = message.properties.get('reply_to')
        try:
            self.current_task = task
            self.current_args = args
            self.current_kwargs = kwargs
            stop_heartbeat = threading.Event()
            heartbeat_thread = threading.Thread(
                target=self._heartbeat_tick,
                args=(message.channel.connection, stop_heartbeat))
            heartbeat_thread.start()
            try:
                result = self._apply_task(task, args, kwargs)
            finally:
                self.current_task = None
                self.current_args = None
                self.current_kwargs = None
                stop_heartbeat.set()
                heartbeat_thread.join()
        except Reject:
            logger.warning('Task is rejected')
            time.sleep(1)  # Prevent cpu burning
            message.channel.basic_reject(message.delivery_tag, requeue=True)
        except Discard:
            logger.warning('Task is discarded')
            message.channel.basic_reject(message.delivery_tag, requeue=False)
        except Exception:
            logger.error('Task raised an exception')
            exc_info = sys.exc_info()
            logger.error(''.join(traceback.format_exception(*exc_info)))
            signals.worker_failure.send(self.kuyruk, description=description,
                                        task=task, args=args, kwargs=kwargs,
                                        exc_info=exc_info, worker=self)
            message.channel.basic_reject(message.delivery_tag, requeue=False)
            if reply_to:
                self._send_reply(reply_to, message.channel, None, exc_info)
        else:
            logger.info('Task is successful')
            message.channel.basic_ack(message.delivery_tag)
            if reply_to:
                self._send_reply(reply_to, message.channel, result, None)
        finally:
            logger.debug("Task is processed")

    @staticmethod
    def _apply_task(task, args, kwargs):
        """Logs the time spent while running the task."""
        if args is None:
            args = ()
        if kwargs is None:
            kwargs = {}

        start = time.time()
        try:
            return task.apply(*args, **kwargs)
        finally:
            end = time.time()
            logger.info("%s finished in %i seconds." % (task.name, end - start))

    def _send_reply(self, reply_to, channel, result, exc_info):
        logger.debug("Sending reply result=%r", result)
        msg = {
            'result': result,
        }
        if exc_info:
            msg['exception'] = {
                'type': str(exc_info[0].__name__),
                'value': str(exc_info[1]),
                'traceback': ''.join(traceback.format_exception(*exc_info)),
            }
        msg = amqp.Message(body=json.dumps(msg))
        channel.basic_publish(msg, exchange="", routing_key=reply_to)

    def _watch_load(self):
        """Pause consuming messages if lood goes above the allowed limit."""
        while not self.shutdown_pending.is_set():
            load = os.getloadavg()[0]
            if load > self.config.WORKER_MAX_LOAD:
                if self._pause is False:
                    logger.warning(
                        'Load is above the treshold (%.2f/%s), '
                        'pausing consumer', load, self.config.WORKER_MAX_LOAD)
                    self._pause = True
            else:
                if self._pause is True:
                    logger.warning(
                        'Load is below the treshold (%.2f/%s), '
                        'resuming consumer', load, self.config.WORKER_MAX_LOAD)
                    self._pause = False
            time.sleep(1)

    @property
    def uptime(self):
        if self._started is not None:
            return os.times()[4] - self._started

    def _shutdown_timer(self):
        """Counts down from MAX_WORKER_RUN_TIME. When it reaches zero sutdown
        gracefully.

        """
        if not self.config.WORKER_MAX_RUN_TIME:
            return

        while not self.shutdown_pending.is_set():
            remaining = self.config.WORKER_MAX_RUN_TIME - self.uptime
            if remaining > 0:
                time.sleep(remaining)
            else:
                logger.warning('Run time reached zero')
                self.shutdown()
                break

    def shutdown(self):
        """Exits after the current task is finished."""
        logger.warning("Shutdown requested")
        self.shutdown_pending.set()

    def _handle_sigint(self, signum, frame):
        """Shutdown after processing current task."""
        logger.warning("Catched SIGINT")
        self.shutdown()

    def _handle_sigterm(self, signum, frame):
        """Shutdown after processing current task."""
        logger.warning("Catched SIGTERM")
        self.shutdown()

    def _handle_sighup(self, signum, frame):
        """Used internally to fail the task when connection to RabbitMQ is
        lost during the execution of the task.

        """
        logger.warning("Catched SIGHUP")
        raise ConnectionError(self._heartbeat_exc_info)

    @staticmethod
    def _handle_sigusr1(signum, frame):
        """Print stacktrace."""
        print('=' * 70)
        print(''.join(traceback.format_stack()))
        print('-' * 70)

    def _handle_sigusr2(self, signum, frame):
        """Drop current task."""
        logger.warning("Catched SIGUSR2")
        if self.current_task:
            logger.warning("Dropping current task...")
            raise Discard

    def drop_task(self):
        os.kill(os.getpid(), signal.SIGUSR2)

    def _heartbeat_tick(self, connection, stop_event):
        while not stop_event.wait(1):
            try:
                connection.heartbeat_tick()
            except socket.timeout:
                pass
            except Exception as e:
                logger.error(e)
                self._heartbeat_exc_info = sys.exc_info()
                os.kill(os.getpid(), signal.SIGHUP)
                break
<commit_msg>handle EINTR on heartbeat thread<commit_after>from __future__ import absolute_import, print_function
import os
import sys
import json
import time
import errno
import socket
import signal
import logging
import logging.config
import warnings
import threading
import traceback
import multiprocessing

import amqp

from kuyruk import importer, signals
from kuyruk.task import get_queue_name
from kuyruk.exceptions import Reject, Discard, ConnectionError

logger = logging.getLogger(__name__)

class Worker(object):
    """Consumes tasks from a queue and runs them.

    :param app: An instance of :class:`~kuyruk.Kuyruk`
    :param args: Command line arguments

    """
    def __init__(self, app, args):
        self.kuyruk = app

        if not args.queues:
            raise ValueError("no queue given")

        self.queues = [get_queue_name(q, local=args.local) for q in args.queues]
        self.shutdown_pending = threading.Event()
        self._pause = False
        self._started = None
        self.consuming = False
        self.current_task = None
        self.current_args = None
        self.current_kwargs = None
        self._heartbeat_exc_info = None
        if self.config.WORKER_MAX_LOAD is None:
            self.config.WORKER_MAX_LOAD = multiprocessing.cpu_count()

        self._pid = os.getpid()
        self._hostname = socket.gethostname()

        signals.worker_init.send(self.kuyruk, worker=self)

    @property
    def queue(self):
        warnings.warn("Worker.queue is deprecated. Use Worker.queues instead.")
        return self.queues[0]

    @property
    def config(self):
        return self.kuyruk.config

    def run(self):
        """Runs the worker and consumes messages from RabbitMQ.
        Returns only after `shutdown()` is called.

        """
        # Lazy import setproctitle.
        # There is bug with the latest version of Python with
        # uWSGI and setproctitle combination.
        # Watch: https://github.com/unbit/uwsgi/issues/1030
        from setproctitle import setproctitle
        setproctitle("kuyruk: worker on %s" % ','.join(self.queues))

        self._setup_logging()

        signal.signal(signal.SIGINT, self._handle_sigint)
        signal.signal(signal.SIGTERM, self._handle_sigterm)
        signal.signal(signal.SIGHUP, self._handle_sighup)
        signal.signal(signal.SIGUSR1, self._handle_sigusr1)
        signal.signal(signal.SIGUSR2, self._handle_sigusr2)

        self._started = os.times()[4]

        for f in (self._watch_load, self._shutdown_timer):
            t = threading.Thread(target=f)
            t.daemon = True
            t.start()

        signals.worker_start.send(self.kuyruk, worker=self)
        self._consume_messages()
        signals.worker_shutdown.send(self.kuyruk, worker=self)

        logger.debug("End run worker")

    def _setup_logging(self):
        if self.config.WORKER_LOGGING_CONFIG:
            logging.config.fileConfig(self.config.WORKER_LOGGING_CONFIG)
        else:
            level = getattr(logging, self.config.WORKER_LOGGING_LEVEL.upper())
            fmt = "%(levelname).1s " \
                  "%(name)s.%(funcName)s:%(lineno)d - %(message)s"
            logging.basicConfig(level=level, format=fmt)

    def _consume_messages(self):
        with self.kuyruk.channel() as ch:
            # Set prefetch count to 1. If we don't set this, RabbitMQ keeps
            # sending messages while we are already working on a message.
            ch.basic_qos(0, 1, True)

            self._declare_queues(ch)
            while not self.shutdown_pending.is_set():
                # Consume or pause
                if self._pause and self.consuming:
                    self._cancel_queues(ch)
                    logger.info('Consumer cancelled')
                    self.consuming = False
                elif not self._pause and not self.consuming:
                    self._consume_queues(ch)
                    logger.info('Consumer started')
                    self.consuming = True

                if self._heartbeat_exc_info:
                    break

                try:
                    ch.connection.heartbeat_tick()
                    ch.connection.drain_events(timeout=1)
                except socket.timeout:
                    pass
                except socket.error as e:
                    if e.errno != errno.EINTR:
                        raise
        logger.debug("End run worker")

    def _consumer_tag(self, queue):
        return "%s:%s@%s" % (queue, self._pid, self._hostname)

    def _declare_queues(self, ch):
        for queue in self.queues:
            logger.debug("queue_declare: %s", queue)
            ch.queue_declare(
                queue=queue, durable=True, auto_delete=False)

    def _consume_queues(self, ch):
        for queue in self.queues:
            logger.debug("basic_consume: %s", queue)
            ch.basic_consume(queue=queue,
                             consumer_tag=self._consumer_tag(queue),
                             callback=self._process_message)

    def _cancel_queues(self, ch):
        for queue in self.queues:
            logger.debug("basic_cancel: %s", queue)
            ch.basic_cancel(self._consumer_tag(queue))

    def _process_message(self, message):
        """Processes the message received from the queue."""
        if self.shutdown_pending.is_set():
            return

        try:
            description = json.loads(message.body)
        except Exception:
            message.channel.basic_reject(message.delivery_tag, requeue=False)
            logger.error("Cannot decode message. Dropping.")
        else:
            logger.info("Processing task: %r", description)
            self._process_description(message, description)

    def _process_description(self, message, description):
        try:
            task = importer.import_object(description['module'],
                                          description['function'])
            args, kwargs = description['args'], description['kwargs']
        except Exception:
            logger.error('Cannot import task')
            exc_info = sys.exc_info()
            signals.worker_failure.send(self.kuyruk, description=description,
                                        exc_info=exc_info, worker=self)
            message.channel.basic_reject(message.delivery_tag, requeue=False)
        else:
            self._process_task(message, description, task, args, kwargs)

    def _process_task(self, message, description, task, args, kwargs):
        reply_to = message.properties.get('reply_to')
        try:
            self.current_task = task
            self.current_args = args
            self.current_kwargs = kwargs
            stop_heartbeat = threading.Event()
            heartbeat_thread = threading.Thread(
                target=self._heartbeat_tick,
                args=(message.channel.connection, stop_heartbeat))
            heartbeat_thread.start()
            try:
                result = self._apply_task(task, args, kwargs)
            finally:
                self.current_task = None
                self.current_args = None
                self.current_kwargs = None
                stop_heartbeat.set()
                heartbeat_thread.join()
        except Reject:
            logger.warning('Task is rejected')
            time.sleep(1)  # Prevent cpu burning
            message.channel.basic_reject(message.delivery_tag, requeue=True)
        except Discard:
            logger.warning('Task is discarded')
            message.channel.basic_reject(message.delivery_tag, requeue=False)
        except Exception:
            logger.error('Task raised an exception')
            exc_info = sys.exc_info()
            logger.error(''.join(traceback.format_exception(*exc_info)))
            signals.worker_failure.send(self.kuyruk, description=description,
                                        task=task, args=args, kwargs=kwargs,
                                        exc_info=exc_info, worker=self)
            message.channel.basic_reject(message.delivery_tag, requeue=False)
            if reply_to:
                self._send_reply(reply_to, message.channel, None, exc_info)
        else:
            logger.info('Task is successful')
            message.channel.basic_ack(message.delivery_tag)
            if reply_to:
                self._send_reply(reply_to, message.channel, result, None)
        finally:
            logger.debug("Task is processed")

    @staticmethod
    def _apply_task(task, args, kwargs):
        """Logs the time spent while running the task."""
        if args is None:
            args = ()
        if kwargs is None:
            kwargs = {}

        start = time.time()
        try:
            return task.apply(*args, **kwargs)
        finally:
            end = time.time()
            logger.info("%s finished in %i seconds." % (task.name, end - start))

    def _send_reply(self, reply_to, channel, result, exc_info):
        logger.debug("Sending reply result=%r", result)
        msg = {
            'result': result,
        }
        if exc_info:
            msg['exception'] = {
                'type': str(exc_info[0].__name__),
                'value': str(exc_info[1]),
                'traceback': ''.join(traceback.format_exception(*exc_info)),
            }
        msg = amqp.Message(body=json.dumps(msg))
        channel.basic_publish(msg, exchange="", routing_key=reply_to)

    def _watch_load(self):
        """Pause consuming messages if lood goes above the allowed limit."""
        while not self.shutdown_pending.is_set():
            load = os.getloadavg()[0]
            if load > self.config.WORKER_MAX_LOAD:
                if self._pause is False:
                    logger.warning(
                        'Load is above the treshold (%.2f/%s), '
                        'pausing consumer', load, self.config.WORKER_MAX_LOAD)
                    self._pause = True
            else:
                if self._pause is True:
                    logger.warning(
                        'Load is below the treshold (%.2f/%s), '
                        'resuming consumer', load, self.config.WORKER_MAX_LOAD)
                    self._pause = False
            time.sleep(1)

    @property
    def uptime(self):
        if self._started is not None:
            return os.times()[4] - self._started

    def _shutdown_timer(self):
        """Counts down from MAX_WORKER_RUN_TIME. When it reaches zero sutdown
        gracefully.

        """
        if not self.config.WORKER_MAX_RUN_TIME:
            return

        while not self.shutdown_pending.is_set():
            remaining = self.config.WORKER_MAX_RUN_TIME - self.uptime
            if remaining > 0:
                time.sleep(remaining)
            else:
                logger.warning('Run time reached zero')
                self.shutdown()
                break

    def shutdown(self):
        """Exits after the current task is finished."""
        logger.warning("Shutdown requested")
        self.shutdown_pending.set()

    def _handle_sigint(self, signum, frame):
        """Shutdown after processing current task."""
        logger.warning("Catched SIGINT")
        self.shutdown()

    def _handle_sigterm(self, signum, frame):
        """Shutdown after processing current task."""
        logger.warning("Catched SIGTERM")
        self.shutdown()

    def _handle_sighup(self, signum, frame):
        """Used internally to fail the task when connection to RabbitMQ is
        lost during the execution of the task.

        """
        logger.warning("Catched SIGHUP")
        raise ConnectionError(self._heartbeat_exc_info)

    @staticmethod
    def _handle_sigusr1(signum, frame):
        """Print stacktrace."""
        print('=' * 70)
        print(''.join(traceback.format_stack()))
        print('-' * 70)

    def _handle_sigusr2(self, signum, frame):
        """Drop current task."""
        logger.warning("Catched SIGUSR2")
        if self.current_task:
            logger.warning("Dropping current task...")
            raise Discard

    def drop_task(self):
        os.kill(os.getpid(), signal.SIGUSR2)

    def _heartbeat_tick(self, connection, stop_event):
        while not stop_event.wait(1):
            try:
                try:
                    connection.heartbeat_tick()
                except socket.error as e:
                    if e.errno != errno.EINTR:
                        raise
            except socket.timeout:
                pass
            except Exception as e:
                logger.error(e)
                self._heartbeat_exc_info = sys.exc_info()
                os.kill(os.getpid(), signal.SIGHUP)
                break
<|endoftext|>


<commit_before>#!/usr/bin/env python3

# In this sample you will see how the access control firewall rules of
# each server host can be retrieved and updated using the REST API.
# Access control firewall rules are the custom iptable rules defined by
# the users and are different from the default iptable rules used by the
# QRadar system. Updating the access control firewall rules won't affact
# the default rules in QRadar deployment. The example shows how to add a
# new rule to open a new port on each managed host.

import json
import os
import sys

import importlib
sys.path.append(os.path.realpath('../modules'))
client_module = importlib.import_module('RestApiClient')
SampleUtilities = importlib.import_module('SampleUtilities')

def main():

    # First we have to create our client
    client = client_module.RestApiClient(version='5.0')

    # -------------------------------------------------------------------------
    # Basic 'GET'
    # In this example we'll be using the GET endpoint of system/servers without
    # any parameters. This will print all the servers in the deployment

    # Send in the request
    SampleUtilities.pretty_print_request(client, 'system/servers', 'GET')
    response = client.call_api('system/servers', 'GET')

    # Check if the success code was returned to ensure the call to the API was
    # successful.
    if (response.code != 200):
        print('Failed to retrieve the list of servers')
        SampleUtilities.pretty_print_response(response)
        sys.exit(1)

    # loop through the servers
    response_body = json.loads(response.read().decode('utf-8'))
    print('Number of servers retrieved: ' + str(len(response_body)))

    for server in response_body:
        server_id = server['server_id']
        print('server_id={0}'.format(server_id))

        print('Get the current firewall rules.')
        endpoint = '/system/servers/' + str(server_id) + '/firewall_rules'
        SampleUtilities.pretty_print_request(client, endpoint, 'GET')
        response = client.call_api(endpoint, 'GET')

        if (response.code != 200):
            print('Failed to retrieve the firewall rules of the server ' + str(
                server_id))
            SampleUtilities.pretty_print_response(response)
            sys.exit(1)

        response_body = json.loads(response.read().decode('utf-8'))
        print(json.dumps(response_body, indent=4))

        print('Add a new firewall rule and save back to the server.')
        # Have the user input the server Id.
        ip = input(
            'Please enter the ip of the host from where to ' +
            'access the server.\n')
        port = input(
            'Please enter the port on the server you want to access.\n')
        protocol = input(
            'Please enter the protocol (TCP, UDP, ANY) used for the access.\n')
        response_body.append({
            "is_any_source_ip": True,
            "port_range": None,
            "port_type": "SINGLE",
            "protocol": protocol,
            "single_port": port,
            "source_ip": ip
        })
        headers = {}
        headers['Content-Type'] = 'application/json'
        endpoint = 'system/servers/' + str(server_id) + '/firewall_rules'
        SampleUtilities.pretty_print_request(
            client,
            endpoint,
            'PUT',
            headers=headers)
        response = client.call_api(
            endpoint,
            'PUT',
            headers=headers,
            data=json.dumps(response_body).encode())
        SampleUtilities.pretty_print_response(response)
        if (response.code != 200):
            print('Failed to update the firewall rules of the server ' + str(
                server_id))
            sys.exit(1)

if __name__ == "__main__":
    main()
<commit_msg>Re RTC#97094 Correct the is_any_source_ip flag.<commit_after>#!/usr/bin/env python3

# In this sample you will see how the access control firewall rules of
# each server host can be retrieved and updated using the REST API.
# Access control firewall rules are the custom iptable rules defined by
# the users and are different from the default iptable rules used by the
# QRadar system. Updating the access control firewall rules won't affact
# the default rules in QRadar deployment. The example shows how to add a
# new rule to open a new port on each managed host.

import json
import os
import sys

import importlib
sys.path.append(os.path.realpath('../modules'))
client_module = importlib.import_module('RestApiClient')
SampleUtilities = importlib.import_module('SampleUtilities')

def main():

    # First we have to create our client
    client = client_module.RestApiClient(version='5.0')

    # -------------------------------------------------------------------------
    # Basic 'GET'
    # In this example we'll be using the GET endpoint of system/servers without
    # any parameters. This will print all the servers in the deployment

    # Send in the request
    SampleUtilities.pretty_print_request(client, 'system/servers', 'GET')
    response = client.call_api('system/servers', 'GET')

    # Check if the success code was returned to ensure the call to the API was
    # successful.
    if (response.code != 200):
        print('Failed to retrieve the list of servers')
        SampleUtilities.pretty_print_response(response)
        sys.exit(1)

    # loop through the servers
    response_body = json.loads(response.read().decode('utf-8'))
    print('Number of servers retrieved: ' + str(len(response_body)))

    for server in response_body:
        server_id = server['server_id']
        print('server_id={0}'.format(server_id))

        print('Get the current firewall rules.')
        endpoint = '/system/servers/' + str(server_id) + '/firewall_rules'
        SampleUtilities.pretty_print_request(client, endpoint, 'GET')
        response = client.call_api(endpoint, 'GET')

        if (response.code != 200):
            print('Failed to retrieve the firewall rules of the server ' + str(
                server_id))
            SampleUtilities.pretty_print_response(response)
            sys.exit(1)

        response_body = json.loads(response.read().decode('utf-8'))
        print(json.dumps(response_body, indent=4))

        print('Add a new firewall rule and save back to the server.')
        # Have the user input the server Id.
        ip = input(
            'Please enter the ip of the host from where to ' +
            'access the server.\n')
        port = input(
            'Please enter the port on the server you want to access.\n')
        protocol = input(
            'Please enter the protocol (TCP, UDP, ANY) used for the access.\n')
        response_body.append({
            "is_any_source_ip": False,
            "port_range": None,
            "port_type": "SINGLE",
            "protocol": protocol,
            "single_port": port,
            "source_ip": ip
        })
        headers = {}
        headers['Content-Type'] = 'application/json'
        endpoint = 'system/servers/' + str(server_id) + '/firewall_rules'
        SampleUtilities.pretty_print_request(
            client,
            endpoint,
            'PUT',
            headers=headers)
        response = client.call_api(
            endpoint,
            'PUT',
            headers=headers,
            data=json.dumps(response_body).encode())
        SampleUtilities.pretty_print_response(response)
        if (response.code != 200):
            print('Failed to update the firewall rules of the server ' + str(
                server_id))
            sys.exit(1)

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