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

# FITS Viewer

# The MIT License
#
# Copyright (c) 2016 Jeremie DECOCK (http://www.jdhp.org)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# 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 OR COPYRIGHT HOLDERS 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.

# Documentation: http://docs.astropy.org/en/stable/io/fits/index.html

__all__ = ['TkGUI']

import matplotlib
matplotlib.use('TkAgg')

import matplotlib.pyplot as plt

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2TkAgg
# implement the default mpl key bindings
from matplotlib.backend_bases import key_press_handler
from matplotlib import cm

import tkinter as tk

import argparse
import os

from astropy.io import fits

###############################################################################

def get_image_array_from_fits_file(file_path):

    hdu_list = fits.open(file_path)   # open the FITS file

    if len(hdu_list) != 1:
        raise Exception("The FITS file should contain only one HDU.")

    image_array = hdu_list[0].data    # "hdu.data" is a Numpy Array

    hdu_list.close()

    return image_array

def get_colour_map_list():
    """Return the list of the available colormaps."""
    return sorted(plt.cm.datad)

###############################################################################

class TkGUI:
    """
    TODO...
    """

    def __init__(self, root):
        """
        TODO...
        """

        # Matplotlib ##################

        self.fig = plt.figure(figsize=(8.0, 8.0))
        self.ax = self.fig.add_subplot(111)

        # Gui parameters ##############

        self.color_map = None

        self._show_color_bar = tk.BooleanVar()

        self.file_path = None
        self.image_array = None

        # Make widgets ################

        #self.root = tk.Tk()   # TODO
        self.root = root

        # Add a callback on WM_DELETE_WINDOW events
        self.root.protocol("WM_DELETE_WINDOW", self.quit)

        # Canvas
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.get_tk_widget().pack(fill="both", expand=True)

        # Buttons
        quit_button = tk.Button(master=self.root, text='Quit', command=self.quit)
        quit_button.pack(fill="x", expand=True)

        # Make a menubar ##############

        # Create a toplevel menu
        menubar = tk.Menu(self.root)

        # Create a pulldown menu
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="Open...", command=self.select_fits_file)
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self.quit)

        menubar.add_cascade(label="File", menu=file_menu)

        # Create a pulldown menu
        view_menu = tk.Menu(menubar, tearoff=0)
        view_menu.add_checkbutton(label="Show color bar",
                                  variable=self._show_color_bar,
                                  command=self.draw_figure)

        menubar.add_cascade(label="View", menu=view_menu)

        # Display the menu
        # The config method is used to attach the menu to the root window. The
        # contents of that menu is used to create a menubar at the top of the root
        # window. There is no need to pack the menu, since it is automatically
        # displayed by Tkinter.
        self.root.config(menu=menubar)

    def run(self):
        """
        Launch the main loop (Tk event loop).
        """

        # TODO ???
        self.root.mainloop()

    def quit(self):
        self.root.quit()     # stops mainloop
        self.root.destroy()  # this is necessary on Windows to prevent
                             # Fatal Python Error: PyEval_RestoreThread: NULL tstate

    def select_fits_file(self):
        """
        Display a file dialog to select the FITS file to open.

        Return the path of the selected file.
        """

        # Check and parse the file

        # FILE_TYPES = [(label1, pattern1), (label2, pattern2), ...]
        FILE_TYPES = [
                    ('FITS Files', '.fits')
                ]

        HOME = os.path.expanduser("~")

        path = tk.filedialog.askopenfilename(parent=self.root,
                                             filetypes=FILE_TYPES,     # optional
                                             defaultextension='.fits', # optional
                                             initialdir=HOME,          # optional
                                             #initialfile='demo.fits',  # optional
                                             title='Select your file') # optional

        self.open_fits_file(path)

    def open_fits_file(self, file_path):
        """
        Open and display the given FITS file.
        """
        # READ THE INPUT FILE #################################################

        self.file_path = file_path
        self.image_array = get_image_array_from_fits_file(self.file_path)

        if self.image_array.ndim != 2:
            raise Exception("Unexpected error: the input FITS file should contain a 2D array.")

        self.root.title(self.file_path)
        self.draw_figure()

    def close_fits_file(self):
        self.file_path = None
        self.image_array = None
        self.clear_figure()

    def clear_figure(self):
        self.fig.clf()

    def draw_figure(self):
        self.fig.clf() # TODO
        self.ax = self.fig.add_subplot(111)

        #self.ax.set_title(self.file_path)

        im = self.ax.imshow(self.image_array,
                            origin='lower',
                            interpolation='nearest',
                            cmap=self.color_map)

        if self.show_color_bar:
            plt.colorbar(im) # draw the colorbar

        self.fig.canvas.draw()

    # PROPERTIES ##############################################################

    @property
    def show_color_bar(self):
        return self._show_color_bar.get()

    @show_color_bar.setter
    def show_color_bar(self, value):
        self._show_color_bar.set(value)

def main():

    root = tk.Tk()   # TODO
    gui = TkGUI(root)

    # PARSE OPTIONS ###########################################################

    parser = argparse.ArgumentParser(description="Display a FITS file.")

    parser.add_argument("--cmap", "-c", default="gray", metavar="STRING",
            help="the colormap to use. The list of available color maps is available here: "
                 "http://matplotlib.org/examples/color/colormaps_reference.html")

    parser.add_argument("--hidecbar", "-H", action="store_true",
            help="hide the color bar")

    parser.add_argument("filearg", nargs=1, metavar="FILE",
            help="the FITS file to process")

    args = parser.parse_args()

    input_file_path = args.filearg[0]

    # SET OPTIONS #############################################################

    gui.color_map = args.cmap
    gui.show_color_bar = not args.hidecbar

    gui.open_fits_file(input_file_path)

    # LAUNCH THE MAIN LOOP ####################################################

    gui.run()

if __name__ == "__main__":
    main()

<commit_msg>Add a 'close' button + misc updates.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# FITS Viewer

# The MIT License
#
# Copyright (c) 2016 Jeremie DECOCK (http://www.jdhp.org)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# 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 OR COPYRIGHT HOLDERS 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.

# Documentation: http://docs.astropy.org/en/stable/io/fits/index.html

__all__ = ['TkGUI']

import matplotlib
matplotlib.use('TkAgg')

import matplotlib.pyplot as plt

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2TkAgg
# implement the default mpl key bindings
from matplotlib.backend_bases import key_press_handler
from matplotlib import cm

import tkinter as tk

import argparse
import os

from astropy.io import fits

###############################################################################

def get_image_array_from_fits_file(file_path):

    hdu_list = fits.open(file_path)   # open the FITS file

    if len(hdu_list) != 1:
        raise Exception("The FITS file should contain only one HDU.")

    image_array = hdu_list[0].data    # "hdu.data" is a Numpy Array

    hdu_list.close()

    return image_array

def get_colour_map_list():
    """Return the list of the available colormaps."""
    return sorted(plt.cm.datad)

###############################################################################

class TkGUI:
    """
    TODO...
    """

    def __init__(self, root):
        """
        TODO...
        """

        self.file_path = None
        self.image_array = None

        # Matplotlib ##################

        self.fig = plt.figure(figsize=(8.0, 8.0))
        self.ax = self.fig.add_subplot(111)

        # Gui parameters ##############

        self.color_map = None
        self._show_color_bar = tk.BooleanVar()

        # Make widgets ################

        self.root = root

        # Add a callback on WM_DELETE_WINDOW events
        self.root.protocol("WM_DELETE_WINDOW", self.quit)

        # Canvas
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
        self.canvas.get_tk_widget().pack(fill="both", expand=True)

        # Buttons
        quit_button = tk.Button(master=self.root, text='Quit', command=self.quit)
        quit_button.pack(fill="x", expand=True)

        # Make a menubar ##############

        # Create a toplevel menu
        menubar = tk.Menu(self.root)

        # Create a pulldown menu
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="Open...", command=self.select_fits_file)
        file_menu.add_command(label="Close", command=self.close_fits_file)
        file_menu.add_separator()
        file_menu.add_command(label="Exit", command=self.quit)

        menubar.add_cascade(label="File", menu=file_menu)

        # Create a pulldown menu
        view_menu = tk.Menu(menubar, tearoff=0)
        view_menu.add_checkbutton(label="Show color bar",
                                  variable=self._show_color_bar,
                                  command=self.draw_figure)

        menubar.add_cascade(label="View", menu=view_menu)

        # Display the menu
        # The config method is used to attach the menu to the root window. The
        # contents of that menu is used to create a menubar at the top of the root
        # window. There is no need to pack the menu, since it is automatically
        # displayed by Tkinter.
        self.root.config(menu=menubar)

    def run(self):
        """
        Launch the main loop (Tk event loop).
        """

        # TODO ???
        self.root.mainloop()

    def quit(self):
        self.root.quit()     # stops mainloop
        self.root.destroy()  # this is necessary on Windows to prevent
                             # Fatal Python Error: PyEval_RestoreThread: NULL tstate

    def select_fits_file(self):
        """
        Display a file dialog to select the FITS file to open.

        Return the path of the selected file.
        """

        # Check and parse the file

        # FILE_TYPES = [(label1, pattern1), (label2, pattern2), ...]
        FILE_TYPES = [
                    ('FITS Files', '.fits')
                ]

        HOME = os.path.expanduser("~")

        path = tk.filedialog.askopenfilename(parent=self.root,
                                             filetypes=FILE_TYPES,     # optional
                                             defaultextension='.fits', # optional
                                             initialdir=HOME,          # optional
                                             #initialfile='demo.fits',  # optional
                                             title='Select your file') # optional

        self.open_fits_file(path)

    def open_fits_file(self, file_path):
        """
        Open and display the given FITS file.
        """
        # READ THE INPUT FILE #################################################

        self.file_path = file_path
        self.image_array = get_image_array_from_fits_file(self.file_path)

        if self.image_array.ndim != 2:
            raise Exception("Unexpected error: the input FITS file should contain a 2D array.")

        self.root.title(self.file_path)
        self.draw_figure()

    def close_fits_file(self):
        self.file_path = None
        self.image_array = None
        self.clear_figure()

    def clear_figure(self):
        self.fig.clf()
        self.fig.canvas.draw()

    def draw_figure(self):
        if self.image_array is not None:
            self.fig.clf() # TODO
            self.ax = self.fig.add_subplot(111)

            #self.ax.set_title(self.file_path)

            im = self.ax.imshow(self.image_array,
                                origin='lower',
                                interpolation='nearest',
                                cmap=self.color_map)

            if self.show_color_bar:
                plt.colorbar(im) # draw the colorbar

            self.fig.canvas.draw()

    # PROPERTIES ##############################################################

    @property
    def show_color_bar(self):
        return self._show_color_bar.get()

    @show_color_bar.setter
    def show_color_bar(self, value):
        self._show_color_bar.set(value)

def main():

    root = tk.Tk()   # TODO
    gui = TkGUI(root)

    # PARSE OPTIONS ###########################################################

    parser = argparse.ArgumentParser(description="Display a FITS file.")

    parser.add_argument("--cmap", "-c", default="gray", metavar="STRING",
            help="the colormap to use. The list of available color maps is available here: "
                 "http://matplotlib.org/examples/color/colormaps_reference.html")

    parser.add_argument("--hidecbar", "-H", action="store_true",
            help="hide the color bar")

    parser.add_argument("filearg", nargs=1, metavar="FILE",
            help="the FITS file to process")

    args = parser.parse_args()

    input_file_path = args.filearg[0]

    # SET OPTIONS #############################################################

    gui.color_map = args.cmap
    gui.show_color_bar = not args.hidecbar

    gui.open_fits_file(input_file_path)

    # LAUNCH THE MAIN LOOP ####################################################

    gui.run()

if __name__ == "__main__":
    main()

<|endoftext|>


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

from __future__ import absolute_import, print_function, unicode_literals
from fsed.utils import open_file
import click
import fsed.ahocorasick
import logging
import re
import sys

logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s',
                    stream=sys.stderr, level=logging.INFO)
LOGGER = logging.getLogger(__name__)

def set_log_level(verbose, quiet):
    '''
    Ses the logging level of the script based on command line options.

    Arguments:
    - `verbose`:
    - `quiet`:
    '''
    '''Sets the logging level.'''
    if quiet:
        verbose = -1
    if verbose < 0:
        verbose = logging.CRITICAL
    elif verbose == 1:
        verbose = logging.INFO
    elif 1 < verbose:
        verbose = logging.DEBUG
    LOGGER.setLevel(verbose)
    fsed.ahocorasick.LOGGER.setLevel(verbose)

def detect_pattern_format(pattern_filename, encoding, on_word_boundaries):
    '''
    Automatically detects the pattern file format, and determines
    whether the Aho-Corasick string matching should pay attention to
    word boundaries or not.

    Arguments:
    - `pattern_filename`:
    - `encoding`:
    - `on_word_boundaries`:
    '''
    tsv = True
    boundaries = on_word_boundaries
    with open_file(pattern_filename) as input_file:
        for line in input_file:
            line = line.decode(encoding)
            if line.count('\t') != 1:
                tsv = False
            if '\\b' in line:
                boundaries = True
            if boundaries and not tsv:
                break
    return tsv, boundaries

def sub_escapes(sval):
    '''
    Process escaped characters in ``sval``.

    Arguments:
    - `sval`:
    '''
    sval = sval.replace('\\a', '\a')
    sval = sval.replace('\\b', '\x00')
    sval = sval.replace('\\f', '\f')
    sval = sval.replace('\\n', '\n')
    sval = sval.replace('\\r', '\r')
    sval = sval.replace('\\t', '\t')
    sval = sval.replace('\\v', '\v')
    sval = sval.replace('\\\\', '\\')
    return sval

def build_trie(pattern_filename, pattern_format, encoding, on_word_boundaries):
    '''
    Constructs a finite state machine for performing string rewriting.

    Arguments:
    - `pattern_filename`:
    - `pattern_format`:
    - `encoding`:
    - `on_word_boundaries`:
    '''
    boundaries = on_word_boundaries
    if pattern_format == 'auto' or not on_word_boundaries:
        tsv, boundaries = detect_pattern_format(pattern_filename, encoding,
                                                on_word_boundaries)
    if pattern_format == 'auto':
        if tsv:
            pattern_format = 'tsv'
        else:
            pattern_format = 'sed'
    trie = fsed.ahocorasick.AhoCorasickTrie()
    num_candidates = 0
    with open_file(pattern_filename) as pattern_file:
        for lineno, line in enumerate(pattern_file):
            line = line.decode(encoding).rstrip('\n')
            if not line.strip():
                continue
            # decode the line
            if pattern_format == 'tsv':
                fields = line.split('\t')
                if len(fields) != 2:
                    LOGGER.warning(('skipping line {} of pattern file (not '
                                    'in tab-separated format): {}').format(lineno, line))
                    continue
                before, after = fields
            elif pattern_format == 'sed':
                before = after = None
                line = line.lstrip()
                if line[0] == 's':
                    delim = line[1]
                    fields = re.split(r'(?<!\\){}'.format(delim), line)
                    if len(fields) == 4:
                        before, after = fields[1], fields[2]
                        before = re.sub(r'(?<!\\)\\{}'.format(delim), delim, before)
                        after = re.sub(r'(?<!\\)\\{}'.format(delim), delim, after)
                if before is None or after is None:
                    LOGGER.warning(('skipping line {} of pattern file (not '
                                    'in sed format): {}').format(lineno, line))
                    continue
            num_candidates += 1
            if on_word_boundaries and before != before.strip():
                LOGGER.warning(('before pattern on line {} padded whitespace; '
                                'this may interact strangely with the --words '
                                'option: {}').format(lineno, line))
            before = sub_escapes(before)
            after = sub_escapes(after)
            if boundaries:
                before = fsed.ahocorasick.boundary_transform(before, on_word_boundaries)
            trie[before] = after
    LOGGER.info('{} patterns loaded from {}'.format(num_candidates,
                                                    pattern_filename))
    return trie, boundaries

def rewrite_str_with_trie(sval, trie, boundaries = False, slow = False):
    '''
    Rewrites a string using the given trie object.

    Arguments:
    - `sval`:
    - `trie`:
    - `boundaries`:
    - `slow`:
    '''
    if boundaries:
        sval = fsed.ahocorasick.boundary_transform(sval)
    if slow:
        sval = trie.replace(sval)
    else:
        sval = trie.greedy_replace(sval)
    if boundaries:
        sval = ''.join(fsed.ahocorasick.boundary_untransform(sval))
    return sval

@click.command()
@click.argument('pattern_filename', type=click.Path(exists=True),
                metavar='PATTERN_FILE')
@click.argument('input_filenames', default='-', nargs=-1,
                type=click.Path(exists=True), metavar='[INPUT_FILES]')
@click.option('--pattern-format', type=click.Choice(['auto', 'tsv', 'sed']),
              default='auto', show_default=True,
              help='Specify the format of PATTERN_FILE')
@click.option('-o', '--output', 'output_filename', type=click.Path(),
              help='Program output is written '
              'to this file. Default is to write '
              'to standard output.')
@click.option('-e', '--encoding', default='utf-8', show_default=True,
              help='The character encoding to use')
@click.option('-w', '--words', is_flag=True,
              help='Match only on word boundaries: '
              'appends "\\b" to the beginning and '
              'end of every pattern in PATTERN_FILE.')
@click.option('--by-line/--across-lines', default=False,
              help='Process the input line by '
              'line or character by character; the default is --across-lines.')
@click.option('--slow', is_flag=True,
              help='Try very hard to '
              'find the longest matches on the input; this is very slow, '
              'and forces --by-line.')
@click.option('-v', '--verbose', default=0, count=True,
              help='Turns on debugging output.')
@click.option('-q', '--quiet', is_flag=True,
              help='Quiet operation, do not emit warnings.')
def main(pattern_filename, input_filenames, pattern_format,
         output_filename,
         encoding, words, by_line, slow, verbose, quiet):
    '''
    Search and replace on INPUT_FILE(s) (or standard input), with
    matching on fixed strings.
    '''
    set_log_level(verbose, quiet)
    if slow:
        by_line = True
    by_line = True # TODO: implement non-line-based rewriting
    # load the patterns
    LOGGER.info('fsed {} input {} output {}'.format(pattern_filename,
                                                    input_filenames,
                                                    output_filename))
    if not input_filenames:
        input_filenames = ('-',)
    if not output_filename:
        output_filename = '-'
    # build trie machine for matching
    trie, boundaries = build_trie(pattern_filename, pattern_format, encoding, words)
    LOGGER.info('writing to {}'.format(output_filename))
    with open_file(output_filename, 'w') as output_file:
        for input_filename in input_filenames:
            # search and replace
            with open_file(input_filename) as input_file:
                LOGGER.info('reading {}'.format(input_filename))
                if by_line:
                    num_lines = 0
                    for line in input_file:
                        line = line.decode(encoding).rstrip('\n')
                        line = rewrite_str_with_trie(line, trie, boundaries, slow)
                        output_file.write((line + '\n').encode(encoding))
                        num_lines += 1
                    LOGGER.info('{} lines written'.format(num_lines))
                else:
                    raise NotImplementedError

if __name__ == '__main__':
    main()
<commit_msg>fsed/fsed: pylint cleanups<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
fsed.py
(c) Will Roberts  12 December, 2015

Main module for the ``fsed`` command line utility.
'''

from __future__ import absolute_import, print_function, unicode_literals
from fsed.utils import open_file
import click
import fsed.ahocorasick
import logging
import re
import sys

logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s',
                    stream=sys.stderr, level=logging.INFO)
LOGGER = logging.getLogger(__name__)

def set_log_level(verbose, quiet):
    '''
    Ses the logging level of the script based on command line options.

    Arguments:
    - `verbose`:
    - `quiet`:
    '''
    if quiet:
        verbose = -1
    if verbose < 0:
        verbose = logging.CRITICAL
    elif verbose == 1:
        verbose = logging.INFO
    elif 1 < verbose:
        verbose = logging.DEBUG
    LOGGER.setLevel(verbose)
    fsed.ahocorasick.LOGGER.setLevel(verbose)

def detect_pattern_format(pattern_filename, encoding, on_word_boundaries):
    '''
    Automatically detects the pattern file format, and determines
    whether the Aho-Corasick string matching should pay attention to
    word boundaries or not.

    Arguments:
    - `pattern_filename`:
    - `encoding`:
    - `on_word_boundaries`:
    '''
    tsv = True
    boundaries = on_word_boundaries
    with open_file(pattern_filename) as input_file:
        for line in input_file:
            line = line.decode(encoding)
            if line.count('\t') != 1:
                tsv = False
            if '\\b' in line:
                boundaries = True
            if boundaries and not tsv:
                break
    return tsv, boundaries

def sub_escapes(sval):
    '''
    Process escaped characters in ``sval``.

    Arguments:
    - `sval`:
    '''
    sval = sval.replace('\\a', '\a')
    sval = sval.replace('\\b', '\x00')
    sval = sval.replace('\\f', '\f')
    sval = sval.replace('\\n', '\n')
    sval = sval.replace('\\r', '\r')
    sval = sval.replace('\\t', '\t')
    sval = sval.replace('\\v', '\v')
    sval = sval.replace('\\\\', '\\')
    return sval

def build_trie(pattern_filename, pattern_format, encoding, on_word_boundaries):
    '''
    Constructs a finite state machine for performing string rewriting.

    Arguments:
    - `pattern_filename`:
    - `pattern_format`:
    - `encoding`:
    - `on_word_boundaries`:
    '''
    boundaries = on_word_boundaries
    if pattern_format == 'auto' or not on_word_boundaries:
        tsv, boundaries = detect_pattern_format(pattern_filename, encoding,
                                                on_word_boundaries)
    if pattern_format == 'auto':
        if tsv:
            pattern_format = 'tsv'
        else:
            pattern_format = 'sed'
    trie = fsed.ahocorasick.AhoCorasickTrie()
    num_candidates = 0
    with open_file(pattern_filename) as pattern_file:
        for lineno, line in enumerate(pattern_file):
            line = line.decode(encoding).rstrip('\n')
            if not line.strip():
                continue
            # decode the line
            if pattern_format == 'tsv':
                fields = line.split('\t')
                if len(fields) != 2:
                    LOGGER.warning(('skipping line {} of pattern file (not '
                                    'in tab-separated format): {}').format(lineno, line))
                    continue
                before, after = fields
            elif pattern_format == 'sed':
                before = after = None
                line = line.lstrip()
                if line[0] == 's':
                    delim = line[1]
                    fields = re.split(r'(?<!\\){}'.format(delim), line)
                    if len(fields) == 4:
                        before, after = fields[1], fields[2]
                        before = re.sub(r'(?<!\\)\\{}'.format(delim), delim, before)
                        after = re.sub(r'(?<!\\)\\{}'.format(delim), delim, after)
                if before is None or after is None:
                    LOGGER.warning(('skipping line {} of pattern file (not '
                                    'in sed format): {}').format(lineno, line))
                    continue
            num_candidates += 1
            if on_word_boundaries and before != before.strip():
                LOGGER.warning(('before pattern on line {} padded whitespace; '
                                'this may interact strangely with the --words '
                                'option: {}').format(lineno, line))
            before = sub_escapes(before)
            after = sub_escapes(after)
            if boundaries:
                before = fsed.ahocorasick.boundary_transform(before, on_word_boundaries)
            trie[before] = after
    LOGGER.info('{} patterns loaded from {}'.format(num_candidates,
                                                    pattern_filename))
    return trie, boundaries

def rewrite_str_with_trie(sval, trie, boundaries = False, slow = False):
    '''
    Rewrites a string using the given trie object.

    Arguments:
    - `sval`:
    - `trie`:
    - `boundaries`:
    - `slow`:
    '''
    if boundaries:
        sval = fsed.ahocorasick.boundary_transform(sval)
    if slow:
        sval = trie.replace(sval)
    else:
        sval = trie.greedy_replace(sval)
    if boundaries:
        sval = ''.join(fsed.ahocorasick.boundary_untransform(sval))
    return sval

@click.command()
@click.argument('pattern_filename', type=click.Path(exists=True),
                metavar='PATTERN_FILE')
@click.argument('input_filenames', default='-', nargs=-1,
                type=click.Path(exists=True), metavar='[INPUT_FILES]')
@click.option('--pattern-format', type=click.Choice(['auto', 'tsv', 'sed']),
              default='auto', show_default=True,
              help='Specify the format of PATTERN_FILE')
@click.option('-o', '--output', 'output_filename', type=click.Path(),
              help='Program output is written '
              'to this file. Default is to write '
              'to standard output.')
@click.option('-e', '--encoding', default='utf-8', show_default=True,
              help='The character encoding to use')
@click.option('-w', '--words', is_flag=True,
              help='Match only on word boundaries: '
              'appends "\\b" to the beginning and '
              'end of every pattern in PATTERN_FILE.')
@click.option('--by-line/--across-lines', default=False,
              help='Process the input line by '
              'line or character by character; the default is --across-lines.')
@click.option('--slow', is_flag=True,
              help='Try very hard to '
              'find the longest matches on the input; this is very slow, '
              'and forces --by-line.')
@click.option('-v', '--verbose', default=0, count=True,
              help='Turns on debugging output.')
@click.option('-q', '--quiet', is_flag=True,
              help='Quiet operation, do not emit warnings.')
def main(pattern_filename, input_filenames, pattern_format,
         output_filename,
         encoding, words, by_line, slow, verbose, quiet):
    '''
    Search and replace on INPUT_FILE(s) (or standard input), with
    matching on fixed strings.
    '''
    set_log_level(verbose, quiet)
    if slow:
        by_line = True
    by_line = True # TODO: implement non-line-based rewriting
    # load the patterns
    LOGGER.info('fsed {} input {} output {}'.format(pattern_filename,
                                                    input_filenames,
                                                    output_filename))
    if not input_filenames:
        input_filenames = ('-',)
    if not output_filename:
        output_filename = '-'
    # build trie machine for matching
    trie, boundaries = build_trie(pattern_filename, pattern_format, encoding, words)
    LOGGER.info('writing to {}'.format(output_filename))
    with open_file(output_filename, 'w') as output_file:
        for input_filename in input_filenames:
            # search and replace
            with open_file(input_filename) as input_file:
                LOGGER.info('reading {}'.format(input_filename))
                if by_line:
                    num_lines = 0
                    for line in input_file:
                        line = line.decode(encoding).rstrip('\n')
                        line = rewrite_str_with_trie(line, trie, boundaries, slow)
                        output_file.write((line + '\n').encode(encoding))
                        num_lines += 1
                    LOGGER.info('{} lines written'.format(num_lines))
                else:
                    raise NotImplementedError

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


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

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

import time
import socket
import threading

import fte.conf
import fte.encoder
import fte.network_io

class worker(threading.Thread):

    """``fte.relay.worker`` is responsible for relaying data between two sockets. Given ``socket1`` and
    ``socket2``, the worker will forward all data
    from ``socket1`` to ``socket2``, and ``socket2`` to ``socket1``. This class is a subclass of
    threading.Thread and does not start relaying until start() is called. The run
    method terminates when either ``socket1`` or ``socket2`` is detected to be closed.
    """

    def __init__(self, socket1, socket2):
        threading.Thread.__init__(self)
        self._socket1 = socket1
        self._socket2 = socket2

    def run(self):
        """It's the responsibility of run to forward data from ``socket1`` to
        ``socket2`` and from ``socket2`` to ``socket1``. The ``run()`` method
        terminates and closes both sockets if ``fte.network_io.recvall_from_socket``
        returns a negative result for ``success``.
        """

        try:
            throttle = fte.conf.getValue('runtime.fte.relay.throttle')
            while True:
                [success, _data] = fte.network_io.recvall_from_socket(
                    self._socket1)
                if not success:
                    break
                if _data:
                    fte.network_io.sendall_to_socket(self._socket2, _data)
                else:
                    time.sleep(throttle)
        except Exception as e:
            fte.warn("fte.worker terminated prematurely" + str(e))
        finally:
            fte.network_io.close_socket(self._socket1)
            fte.network_io.close_socket(self._socket2)

class listener(threading.Thread):

    """It's the responsibility of ``fte.relay.listener`` to bind to
    ``local_ip:local_port``. Once bound it will then relay all incoming connections
    to ``remote_ip:remote_port``.
    All new incoming connections are wrapped with ``onNewIncomingConnection``.
    All new outgoing connections are wrapped with ``onNewOutgoingConnection``.
    By default the functions ``onNewIncomingConnection`` and
    ``onNewOutgoingConnection`` are the identity function.
    """

    def __init__(self, local_ip, local_port,
                 remote_ip, remote_port):
        threading.Thread.__init__(self)

        self._running = False
        self._local_ip = local_ip
        self._local_port = local_port
        self._remote_ip = remote_ip
        self._remote_port = remote_port

    def _instantiateSocket(self):
        try:
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._sock.bind((self._local_ip, self._local_port))
            self._sock.listen(fte.conf.getValue('runtime.fte.relay.backlog'))
            self._sock.settimeout(
                fte.conf.getValue('runtime.fte.relay.accept_timeout'))
        except Exception as e:
            fte.fatal_error('Failed to bind to ' +
                            str((self._local_ip, self._local_port)) + ': ' + str(e))

    def run(self):
        """Bind to ``local_ip:local_port`` and forward all connections to
        ``remote_ip:remote_port``.
        """
        self._instantiateSocket()

        self._running = True
        while self._running:
            try:
                conn, addr = self._sock.accept()

                new_stream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                new_stream.connect((self._remote_ip, self._remote_port))

                conn = self.onNewIncomingConnection(conn)
                new_stream = self.onNewOutgoingConnection(new_stream)

                conn.settimeout(
                    fte.conf.getValue('runtime.fte.relay.socket_timeout'))
                new_stream.settimeout(
                    fte.conf.getValue('runtime.fte.relay.socket_timeout'))

                w1 = worker(conn, new_stream)
                w2 = worker(new_stream, conn)
                w1.start()
                w2.start()
            except socket.timeout:
                continue
            except socket.error:
                fte.warn('socket.error in fte.listener: ' + str(e))
                continue

    def stop(self):
        """Terminate the thread and stop listening on ``local_ip:local_port``.
        """
        self._running = False
        fte.network_io.close_socket(self._sock)

    def onNewIncomingConnection(self, socket):
        """``onNewIncomingConnection`` returns the socket unmodified, by default we do not need to
        perform any modifications to incoming data streams.
        """

        return socket

    def onNewOutgoingConnection(self, socket):
        """``onNewOutgoingConnection`` returns the socket unmodified, by default we do not need to
        perform any modifications to incoming data streams.
        """

        return socket
<commit_msg>Issue #130: catch exceptions in main fte.listener loop<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

import time
import socket
import threading

import fte.conf
import fte.encoder
import fte.network_io

class worker(threading.Thread):

    """``fte.relay.worker`` is responsible for relaying data between two sockets. Given ``socket1`` and
    ``socket2``, the worker will forward all data
    from ``socket1`` to ``socket2``, and ``socket2`` to ``socket1``. This class is a subclass of
    threading.Thread and does not start relaying until start() is called. The run
    method terminates when either ``socket1`` or ``socket2`` is detected to be closed.
    """

    def __init__(self, socket1, socket2):
        threading.Thread.__init__(self)
        self._socket1 = socket1
        self._socket2 = socket2

    def run(self):
        """It's the responsibility of run to forward data from ``socket1`` to
        ``socket2`` and from ``socket2`` to ``socket1``. The ``run()`` method
        terminates and closes both sockets if ``fte.network_io.recvall_from_socket``
        returns a negative result for ``success``.
        """

        try:
            throttle = fte.conf.getValue('runtime.fte.relay.throttle')
            while True:
                [success, _data] = fte.network_io.recvall_from_socket(
                    self._socket1)
                if not success:
                    break
                if _data:
                    fte.network_io.sendall_to_socket(self._socket2, _data)
                else:
                    time.sleep(throttle)
        except Exception as e:
            fte.warn("fte.worker terminated prematurely" + str(e))
        finally:
            fte.network_io.close_socket(self._socket1)
            fte.network_io.close_socket(self._socket2)

class listener(threading.Thread):

    """It's the responsibility of ``fte.relay.listener`` to bind to
    ``local_ip:local_port``. Once bound it will then relay all incoming connections
    to ``remote_ip:remote_port``.
    All new incoming connections are wrapped with ``onNewIncomingConnection``.
    All new outgoing connections are wrapped with ``onNewOutgoingConnection``.
    By default the functions ``onNewIncomingConnection`` and
    ``onNewOutgoingConnection`` are the identity function.
    """

    def __init__(self, local_ip, local_port,
                 remote_ip, remote_port):
        threading.Thread.__init__(self)

        self._running = False
        self._local_ip = local_ip
        self._local_port = local_port
        self._remote_ip = remote_ip
        self._remote_port = remote_port

    def _instantiateSocket(self):
        try:
            self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._sock.bind((self._local_ip, self._local_port))
            self._sock.listen(fte.conf.getValue('runtime.fte.relay.backlog'))
            self._sock.settimeout(
                fte.conf.getValue('runtime.fte.relay.accept_timeout'))
        except Exception as e:
            fte.fatal_error('Failed to bind to ' +
                            str((self._local_ip, self._local_port)) + ': ' + str(e))

    def run(self):
        """Bind to ``local_ip:local_port`` and forward all connections to
        ``remote_ip:remote_port``.
        """
        self._instantiateSocket()

        self._running = True
        while self._running:
            try:
                conn, addr = self._sock.accept()

                new_stream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                new_stream.connect((self._remote_ip, self._remote_port))

                conn = self.onNewIncomingConnection(conn)
                new_stream = self.onNewOutgoingConnection(new_stream)

                conn.settimeout(
                    fte.conf.getValue('runtime.fte.relay.socket_timeout'))
                new_stream.settimeout(
                    fte.conf.getValue('runtime.fte.relay.socket_timeout'))

                w1 = worker(conn, new_stream)
                w2 = worker(new_stream, conn)
                w1.start()
                w2.start()
            except socket.timeout:
                continue
            except socket.error as e:
                continue
            except Exception as e:
                fte.warn('exception in fte.listener: ' + str(e))
                break

    def stop(self):
        """Terminate the thread and stop listening on ``local_ip:local_port``.
        """
        self._running = False
        fte.network_io.close_socket(self._sock)

    def onNewIncomingConnection(self, socket):
        """``onNewIncomingConnection`` returns the socket unmodified, by default we do not need to
        perform any modifications to incoming data streams.
        """

        return socket

    def onNewOutgoingConnection(self, socket):
        """``onNewOutgoingConnection`` returns the socket unmodified, by default we do not need to
        perform any modifications to incoming data streams.
        """

        return socket
<|endoftext|>


<commit_before>import os

import numpy as np
import pandas as pd

from industry_to_days import get_census_mapper

CLEAN2011 = os.path.abspath(
        os.path.join(
            os.path.dirname(__file__),
            '2011_clean.csv'))

def _calc_vacation(row):
    took, days, weeks, months = (row['took_vac'], row['days_vac'],
                                 row['weeks_vac'], row['months_vac'])
    if took in [8, 9] or (days in [998, 999]) or (months in [98, 99]) or (
        weeks in [98, 99]):
        return np.nan
    elif took == 5:
        return 0
    else:
        return days + (7 * weeks) + (30 * months)

def _calc_salary(row):
    amt, unit = row['salary_amt'], row['salary_unit']
    if amt in [0.0, 9999998.0] or unit in [0, 7, 8, 9]:
        return np.nan
    if unit == 3:  # one week
        scalar = 52.0
    elif unit == 4:  # two weeks
        scalar = 26.0
    elif unit == 5:  # one month
        scalar = 12.0
    elif unit == 6:  # one year
        scalar = 1.0
    return scalar * amt

def clean(df):
    # make sex into dummy for is_female
    df['sex'] -= 1
    # figure out total vacation
    df['vacation'] = df.apply(_calc_vacation, axis=1)
    # fix salary to be annual amount
    df['salary'] = df.apply(_calc_salary, axis=1)
    # merge industry data
    df['paid_vacation'] = df.industry.map(get_census_mapper())
    # drop old values
    for col in ['took_vac', 'days_vac', 'weeks_vac', 'months_vac', 'industry',
                'salary_amt', 'salary_unit']:
        df.drop(col, axis=1, inplace=True)
    # write output to a file
    with open(CLEAN2011, 'w+') as csv:
        df.to_csv(path_or_buf=csv)

def do_stats(df):
    pass

def main():
    df = None
    if False and os.path.isfile(CLEAN2011):
        df = pd.io.parsers.read_csv(CLEAN2011)
    else:
        csv_file = os.path.abspath(
            os.path.join(
                os.path.dirname(__file__),
                'psid', '2011er.csv'))
        with open(csv_file) as csv:
            df = pd.io.parsers.read_csv(csv)
    clean(df)
    do_stats(df)

if __name__ == '__main__':
    main()
<commit_msg>allow reading from already-clean csv<commit_after>import os

import numpy as np
import pandas as pd

from industry_to_days import get_census_mapper

CLEAN2011 = os.path.abspath(
        os.path.join(
            os.path.dirname(__file__),
            '2011_clean.csv'))

def _calc_vacation(row):
    took, days, weeks, months = (row['took_vac'], row['days_vac'],
                                 row['weeks_vac'], row['months_vac'])
    if took in [8, 9] or (days in [998, 999]) or (months in [98, 99]) or (
        weeks in [98, 99]):
        return np.nan
    elif took == 5:
        return 0
    else:
        return days + (7 * weeks) + (30 * months)

def _calc_salary(row):
    amt, unit = row['salary_amt'], row['salary_unit']
    if amt in [0.0, 9999998.0] or unit in [0, 7, 8, 9]:
        return np.nan
    if unit == 3:  # one week
        scalar = 52.0
    elif unit == 4:  # two weeks
        scalar = 26.0
    elif unit == 5:  # one month
        scalar = 12.0
    elif unit == 6:  # one year
        scalar = 1.0
    return scalar * amt

def clean(df):
    # make sex into dummy for is_female
    df['sex'] -= 1
    # figure out total vacation
    df['vacation'] = df.apply(_calc_vacation, axis=1)
    # fix salary to be annual amount
    df['salary'] = df.apply(_calc_salary, axis=1)
    # merge industry data
    df['paid_vacation'] = df.industry.map(get_census_mapper())
    # drop old values
    for col in ['took_vac', 'days_vac', 'weeks_vac', 'months_vac', 'industry',
                'salary_amt', 'salary_unit']:
        df.drop(col, axis=1, inplace=True)
    # write output to a file
    with open(CLEAN2011, 'w+') as csv:
        df.to_csv(path_or_buf=csv)

def do_stats(df):
    print df.describe()

def main():
    df = None
    if os.path.isfile(CLEAN2011):
        df = pd.io.parsers.read_csv(CLEAN2011)
    else:
        csv_file = os.path.abspath(
            os.path.join(
                os.path.dirname(__file__),
                'psid', '2011er.csv'))
        with open(csv_file) as csv:
            df = pd.io.parsers.read_csv(csv)
        clean(df)
    do_stats(df)

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


<commit_before>import os
import time
import urllib2

import requests
from flask import Flask, request, json, jsonify, render_template
from sqlalchemy import create_engine
from werkzeug.contrib.cache import SimpleCache

import crud
import logger
import model
import backup
from configs import EnjoliverConfig

ec = EnjoliverConfig()

LOGGER = logger.get_logger(__file__)

app = application = Flask(__name__)
cache = SimpleCache()

application.config["MATCHBOX_URI"] = ec.matchbox_uri
application.config["API_URI"] = ec.api_uri

application.config["MATCHBOX_URLS"] = ec.matchbox_urls

application.config["DB_PATH"] = ec.db_path
application.config["DB_URI"] = ec.db_uri

ignition_journal = application.config["IGNITION_JOURNAL_DIR"] = ec.ignition_journal_dir

application.config["BACKUP_BUCKET_NAME"] = ec.backup_bucket_name
application.config["BACKUP_BUCKET_DIRECTORY"] = ec.backup_bucket_directory
application.config["BACKUP_LOCK_KEY"] = ec.backup_lock_key

engine = None

if __name__ == '__main__' or "gunicorn" in os.getenv("SERVER_SOFTWARE", "foreign"):
    for k, v in ec.__dict__.iteritems():
        LOGGER.info("<config> %s=%s" % (k, v))

    # Start the db
    LOGGER.info("Create engine %s" % application.config["DB_URI"])
    engine = create_engine(application.config["DB_URI"])
    # LOGGER.info("Create model %s" % application.config["DB_URI"])
    # model.Base.metadata.create_all(engine)
    LOGGER.info("Engine with <driver: %s> " % engine.driver)

@application.route("/config", methods=["GET"])
def config():
    return jsonify(ec.__dict__)

@application.route('/', methods=['GET'])
def root():
    """
    Map the API
    :return: available routes
    """
    r = [k.rule for k in application.url_map.iter_rules()]
    r = list(set(r))
    r.sort()
    return jsonify(r)

@application.route('/healthz', methods=['GET'])
def healthz():
    """
    Query all services and return the status
    :return: json
    """
    status = {
        "global": True,
        "flask": True,
        "db": False,
        "matchbox": {k: False for k in application.config["MATCHBOX_URLS"]}
    }
    if app.config["MATCHBOX_URI"] is None:
        application.logger.error("MATCHBOX_URI is None")
    for k in status["matchbox"]:
        try:
            r = requests.get("%s%s" % (app.config["MATCHBOX_URI"], k))
            r.close()
            status["matchbox"][k] = True
        except Exception as e:
            status["matchbox"][k] = False
            status["global"] = False
            LOGGER.error(e)
    try:
        status["db"] = crud.health_check(engine=engine, ts=time.time(), who=request.remote_addr)
    except Exception as e:
        status["global"] = False
        LOGGER.error(e)

    app.logger.debug("%s" % status)
    return json.jsonify(status)

@application.route('/discovery', methods=['POST'])
def discovery():
    if request.content_type != "application/json":
        try:
            r = json.loads(request.data)
        except ValueError:
            app.logger.error("ValueError for %s" % request.data)
            return jsonify(
                {
                    u'boot-info': {},
                    u'lldp': {},
                    u'interfaces': []
                }), 400
    else:
        r = request.get_json()

    while cache.get(application.config["BACKUP_LOCK_KEY"]) is not None:
        app.logger.debug("Cache backup is not None")
        time.sleep(0.1)

    try:
        i = crud.InjectDiscovery(engine=engine,
                                 ignition_journal=ignition_journal,
                                 discovery=r)
        new = i.commit_and_close()
        cache.delete("discovery")
        return jsonify({"total_elt": new[0], "new": new[1]})

    except (KeyError, TypeError):
        return jsonify(
            {
                u'boot-info': {},
                u'lldp': {},
                u'interfaces': []
            }), 406

@application.route('/discovery', methods=['GET'])
def discovery_get():
    key = "discovery"
    all_data = cache.get(key)
    if all_data is None:
        fetch = crud.FetchDiscovery(
            engine=engine,
            ignition_journal=ignition_journal
        )
        all_data = fetch.get_all()
        cache.set(key, all_data, timeout=30)

    return jsonify(all_data)

@application.route('/scheduler', methods=['GET'])
def get_all_schedules():
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    all_sch = fetch.get_schedules()
    fetch.close()

    return jsonify(all_sch)

@application.route('/scheduler/<string:role>', methods=['GET'])
def get_schedule_by_role(role):
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    multi = role.split("&")
    data = fetch.get_roles(*multi)
    fetch.close()

    return jsonify(data)

@application.route('/scheduler/available', methods=['GET'])
def get_available_machine():
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    data = fetch.get_available_machines()
    fetch.close()

    return jsonify(data)

@application.route('/scheduler/ip-list/<string:role>', methods=['GET'])
def get_schedule_role_ip_list(role):
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    ip_list_role = fetch.get_role_ip_list(role)
    fetch.close()

    return jsonify(ip_list_role)

@application.route('/scheduler', methods=['POST'])
def schedule_role():
    if request.content_type != "application/json":
        try:
            r = json.loads(request.get_data())
        except ValueError:
            app.logger.error("ValueError for %s" % request.data)
            return jsonify(
                {
                    u"roles": model.ScheduleRoles.roles,
                    u'selector': {
                        "mac": ""
                    }
                }), 400
    else:
        r = request.get_json()

    inject = crud.InjectSchedule(
        engine=engine,
        data=r)
    try:
        inject.apply_roles()
        cache.delete("schedules")
    finally:
        inject.commit_and_close()

    return jsonify(r)

@application.route('/backup/db', methods=['POST'])
def backup_database():
    return backup.backup_sqlite(cache=cache, application=application)

@application.route('/discovery/interfaces', methods=['GET'])
def discovery_interfaces():
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    interfaces = fetch.get_all_interfaces()

    return jsonify(interfaces)

@application.route('/discovery/ignition-journal/<string:uuid>/<string:boot_id>', methods=['GET'])
def discovery_ignition_journal_by_boot_id(uuid, boot_id):
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal(uuid, boot_id=boot_id)

    return jsonify(lines)

@application.route('/discovery/ignition-journal/<string:uuid>', methods=['GET'])
def discovery_ignition_journal_by_uuid(uuid):
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal(uuid)

    return jsonify(lines)

@application.route('/discovery/ignition-journal', methods=['GET'])
def discovery_ignition_journal_summary():
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal_summary()

    return jsonify(lines)

@application.route('/boot.ipxe', methods=['GET'])
@application.route('/boot.ipxe.0', methods=['GET'])
def boot_ipxe():
    """
    Replace the matchbox/boot.ipxe by insert retry for dhcp and full URL for the chain
    :return: str
    """
    try:
        flask_uri = application.config["API_URI"]
        if flask_uri is None:
            raise AttributeError("API_URI is None")
        app.logger.debug("%s" % flask_uri)

    except Exception as e:
        flask_uri = application.config["MATCHBOX_URI"]
        app.logger.error("<%s %s>: %s" % (e, type(e), e.message))
        app.logger.warning("Fall back to MATCHBOX_URI: %s" % flask_uri)
        if flask_uri is None:
            raise AttributeError("BOTH API_URI and MATCHBOX_URI are None")

    response = \
        "#!ipxe\n" \
        "echo start /boot.ipxe\n" \
        ":retry_dhcp\n" \
        "dhcp || goto retry_dhcp\n" \
        "chain %s/ipxe?" \
        "uuid=${uuid}&" \
        "mac=${net0/mac:hexhyp}&" \
        "domain=${domain}&" \
        "hostname=${hostname}&" \
        "serial=${serial}\n" % flask_uri
    app.logger.debug("%s" % response)
    return response

@application.route("/ignition", methods=["GET"])
def ignition():
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        matchbox_resp = requests.get("%s%s" % (matchbox_uri, request.full_path))
        d = matchbox_resp.content
        matchbox_resp.close()
        return d, matchbox_resp.status_code

    return "matchbox=%s" % matchbox_uri, 403

@application.route("/metadata", methods=["GET"])
def metadata():
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        matchbox_resp = requests.get("%s%s" % (matchbox_uri, request.full_path))
        d = matchbox_resp.content
        matchbox_resp.close()
        return d, matchbox_resp.status_code

    return "matchbox=%s" % matchbox_uri, 403

@app.route('/assets', defaults={'path': ''})
@app.route('/assets/<path:path>')
def assets(path):
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        url = "%s/assets/%s" % (matchbox_uri, path)
        matchbox_resp = requests.get(url)
        d = matchbox_resp.content
        matchbox_resp.close()
        return d, matchbox_resp.status_code

    return "matchbox=%s" % matchbox_uri, 403

@application.route('/ipxe', methods=['GET'])
def ipxe():
    """
    Fetch the matchbox/ipxe?<key>=<value> and insert retry for dhcp
    :return: str
    """
    try:
        matchbox_resp = urllib2.urlopen(
            "%s%s" % (
                app.config["MATCHBOX_URI"],
                request.full_path))
        resp_list = matchbox_resp.readlines()
        matchbox_resp.close()
        if len(resp_list) == 4:
            resp_list.insert(1, "echo start /ipxe\n")
        else:
            app.logger.warning("iPXE response is not coherent")

        response = "".join(resp_list)
        app.logger.debug("%s" % response)
        return response, 200

    except urllib2.URLError:
        return "404", 404

@app.errorhandler(404)
def page_not_found(error):
    return '404\n', 404

@application.route('/ui', methods=['GET'])
def user_interface():
    return render_template("index.html")

@application.route('/ui/view/machine', methods=['GET'])
def user_view_machine():
    key = "discovery"
    all_data = cache.get(key)
    if all_data is None:
        fetch = crud.FetchDiscovery(
            engine=engine,
            ignition_journal=ignition_journal
        )
        all_data = fetch.get_all()
        cache.set(key, all_data, timeout=30)

    res = [["created-date", "updated-date", "uuid", "cidr-boot", "mac-boot"]]
    for i in all_data:
        sub_list = list()
        sub_list.append(i["boot-info"]["created-date"])
        sub_list.append(i["boot-info"]["updated-date"])
        sub_list.append(i["boot-info"]["uuid"])
        for j in i["interfaces"]:
            if j["as_boot"]:
                sub_list.append(j["cidrv4"])
                sub_list.append(j["mac"])
        res.append(sub_list)

    return jsonify(res)

if __name__ == "__main__":
    app.logger.setLevel("DEBUG")
    application.run(debug=True)
<commit_msg>Change the content type<commit_after>import os
import time
import urllib2

import requests
from flask import Flask, request, json, jsonify, render_template, Response
from sqlalchemy import create_engine
from werkzeug.contrib.cache import SimpleCache

import crud
import logger
import model
import backup
from configs import EnjoliverConfig

ec = EnjoliverConfig()

LOGGER = logger.get_logger(__file__)

app = application = Flask(__name__)
cache = SimpleCache()

application.config["MATCHBOX_URI"] = ec.matchbox_uri
application.config["API_URI"] = ec.api_uri

application.config["MATCHBOX_URLS"] = ec.matchbox_urls

application.config["DB_PATH"] = ec.db_path
application.config["DB_URI"] = ec.db_uri

ignition_journal = application.config["IGNITION_JOURNAL_DIR"] = ec.ignition_journal_dir

application.config["BACKUP_BUCKET_NAME"] = ec.backup_bucket_name
application.config["BACKUP_BUCKET_DIRECTORY"] = ec.backup_bucket_directory
application.config["BACKUP_LOCK_KEY"] = ec.backup_lock_key

engine = None

if __name__ == '__main__' or "gunicorn" in os.getenv("SERVER_SOFTWARE", "foreign"):
    for k, v in ec.__dict__.iteritems():
        LOGGER.info("<config> %s=%s" % (k, v))

    # Start the db
    LOGGER.info("Create engine %s" % application.config["DB_URI"])
    engine = create_engine(application.config["DB_URI"])
    # LOGGER.info("Create model %s" % application.config["DB_URI"])
    # model.Base.metadata.create_all(engine)
    LOGGER.info("Engine with <driver: %s> " % engine.driver)

@application.route("/config", methods=["GET"])
def config():
    return jsonify(ec.__dict__)

@application.route('/', methods=['GET'])
def root():
    """
    Map the API
    :return: available routes
    """
    r = [k.rule for k in application.url_map.iter_rules()]
    r = list(set(r))
    r.sort()
    return jsonify(r)

@application.route('/healthz', methods=['GET'])
def healthz():
    """
    Query all services and return the status
    :return: json
    """
    status = {
        "global": True,
        "flask": True,
        "db": False,
        "matchbox": {k: False for k in application.config["MATCHBOX_URLS"]}
    }
    if app.config["MATCHBOX_URI"] is None:
        application.logger.error("MATCHBOX_URI is None")
    for k in status["matchbox"]:
        try:
            r = requests.get("%s%s" % (app.config["MATCHBOX_URI"], k))
            r.close()
            status["matchbox"][k] = True
        except Exception as e:
            status["matchbox"][k] = False
            status["global"] = False
            LOGGER.error(e)
    try:
        status["db"] = crud.health_check(engine=engine, ts=time.time(), who=request.remote_addr)
    except Exception as e:
        status["global"] = False
        LOGGER.error(e)

    app.logger.debug("%s" % status)
    return json.jsonify(status)

@application.route('/discovery', methods=['POST'])
def discovery():
    if request.content_type != "application/json":
        try:
            r = json.loads(request.data)
        except ValueError:
            app.logger.error("ValueError for %s" % request.data)
            return jsonify(
                {
                    u'boot-info': {},
                    u'lldp': {},
                    u'interfaces': []
                }), 400
    else:
        r = request.get_json()

    while cache.get(application.config["BACKUP_LOCK_KEY"]) is not None:
        app.logger.debug("Cache backup is not None")
        time.sleep(0.1)

    try:
        i = crud.InjectDiscovery(engine=engine,
                                 ignition_journal=ignition_journal,
                                 discovery=r)
        new = i.commit_and_close()
        cache.delete("discovery")
        return jsonify({"total_elt": new[0], "new": new[1]})

    except (KeyError, TypeError):
        return jsonify(
            {
                u'boot-info': {},
                u'lldp': {},
                u'interfaces': []
            }), 406

@application.route('/discovery', methods=['GET'])
def discovery_get():
    key = "discovery"
    all_data = cache.get(key)
    if all_data is None:
        fetch = crud.FetchDiscovery(
            engine=engine,
            ignition_journal=ignition_journal
        )
        all_data = fetch.get_all()
        cache.set(key, all_data, timeout=30)

    return jsonify(all_data)

@application.route('/scheduler', methods=['GET'])
def get_all_schedules():
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    all_sch = fetch.get_schedules()
    fetch.close()

    return jsonify(all_sch)

@application.route('/scheduler/<string:role>', methods=['GET'])
def get_schedule_by_role(role):
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    multi = role.split("&")
    data = fetch.get_roles(*multi)
    fetch.close()

    return jsonify(data)

@application.route('/scheduler/available', methods=['GET'])
def get_available_machine():
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    data = fetch.get_available_machines()
    fetch.close()

    return jsonify(data)

@application.route('/scheduler/ip-list/<string:role>', methods=['GET'])
def get_schedule_role_ip_list(role):
    fetch = crud.FetchSchedule(
        engine=engine,
    )
    ip_list_role = fetch.get_role_ip_list(role)
    fetch.close()

    return jsonify(ip_list_role)

@application.route('/scheduler', methods=['POST'])
def schedule_role():
    if request.content_type != "application/json":
        try:
            r = json.loads(request.get_data())
        except ValueError:
            app.logger.error("ValueError for %s" % request.data)
            return jsonify(
                {
                    u"roles": model.ScheduleRoles.roles,
                    u'selector': {
                        "mac": ""
                    }
                }), 400
    else:
        r = request.get_json()

    inject = crud.InjectSchedule(
        engine=engine,
        data=r)
    try:
        inject.apply_roles()
        cache.delete("schedules")
    finally:
        inject.commit_and_close()

    return jsonify(r)

@application.route('/backup/db', methods=['POST'])
def backup_database():
    return backup.backup_sqlite(cache=cache, application=application)

@application.route('/discovery/interfaces', methods=['GET'])
def discovery_interfaces():
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    interfaces = fetch.get_all_interfaces()

    return jsonify(interfaces)

@application.route('/discovery/ignition-journal/<string:uuid>/<string:boot_id>', methods=['GET'])
def discovery_ignition_journal_by_boot_id(uuid, boot_id):
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal(uuid, boot_id=boot_id)

    return jsonify(lines)

@application.route('/discovery/ignition-journal/<string:uuid>', methods=['GET'])
def discovery_ignition_journal_by_uuid(uuid):
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal(uuid)

    return jsonify(lines)

@application.route('/discovery/ignition-journal', methods=['GET'])
def discovery_ignition_journal_summary():
    fetch = crud.FetchDiscovery(engine=engine,
                                ignition_journal=ignition_journal)
    lines = fetch.get_ignition_journal_summary()

    return jsonify(lines)

@application.route('/boot.ipxe', methods=['GET'])
@application.route('/boot.ipxe.0', methods=['GET'])
def boot_ipxe():
    """
    Replace the matchbox/boot.ipxe by insert retry for dhcp and full URL for the chain
    :return: str
    """
    try:
        flask_uri = application.config["API_URI"]
        if flask_uri is None:
            raise AttributeError("API_URI is None")
        app.logger.debug("%s" % flask_uri)

    except Exception as e:
        flask_uri = application.config["MATCHBOX_URI"]
        app.logger.error("<%s %s>: %s" % (e, type(e), e.message))
        app.logger.warning("Fall back to MATCHBOX_URI: %s" % flask_uri)
        if flask_uri is None:
            raise AttributeError("BOTH API_URI and MATCHBOX_URI are None")

    response = \
        "#!ipxe\n" \
        "echo start /boot.ipxe\n" \
        ":retry_dhcp\n" \
        "dhcp || goto retry_dhcp\n" \
        "chain %s/ipxe?" \
        "uuid=${uuid}&" \
        "mac=${net0/mac:hexhyp}&" \
        "domain=${domain}&" \
        "hostname=${hostname}&" \
        "serial=${serial}\n" % flask_uri
    app.logger.debug("%s" % response)
    return response

@application.route("/ignition", methods=["GET"])
def ignition():
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        matchbox_resp = requests.get("%s%s" % (matchbox_uri, request.full_path))
        d = matchbox_resp.content
        matchbox_resp.close()
        return d, matchbox_resp.status_code

    return "matchbox=%s" % matchbox_uri, 403

@application.route("/metadata", methods=["GET"])
def metadata():
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        matchbox_resp = requests.get("%s%s" % (matchbox_uri, request.full_path))
        d = matchbox_resp.content
        matchbox_resp.close()
        return d, matchbox_resp.status_code

    return "matchbox=%s" % matchbox_uri, 403

@app.route('/assets', defaults={'path': ''})
@app.route('/assets/<path:path>')
def assets(path):
    matchbox_uri = application.config.get("MATCHBOX_URI")
    if matchbox_uri:
        url = "%s/assets/%s" % (matchbox_uri, path)
        matchbox_resp = requests.get(url)
        d = matchbox_resp.content
        matchbox_resp.close()
        return Response(response=d, mimetype="application/octet-stream")

    return "matchbox=%s" % matchbox_uri, 403

@application.route('/ipxe', methods=['GET'])
def ipxe():
    """
    Fetch the matchbox/ipxe?<key>=<value> and insert retry for dhcp
    :return: str
    """
    try:
        matchbox_resp = urllib2.urlopen(
            "%s%s" % (
                app.config["MATCHBOX_URI"],
                request.full_path))
        resp_list = matchbox_resp.readlines()
        matchbox_resp.close()
        if len(resp_list) == 4:
            resp_list.insert(1, "echo start /ipxe\n")
        else:
            app.logger.warning("iPXE response is not coherent")

        response = "".join(resp_list)
        app.logger.debug("%s" % response)
        return response, 200

    except urllib2.URLError:
        return "404", 404

@app.errorhandler(404)
def page_not_found(error):
    return '404\n', 404

@application.route('/ui', methods=['GET'])
def user_interface():
    return render_template("index.html")

@application.route('/ui/view/machine', methods=['GET'])
def user_view_machine():
    key = "discovery"
    all_data = cache.get(key)
    if all_data is None:
        fetch = crud.FetchDiscovery(
            engine=engine,
            ignition_journal=ignition_journal
        )
        all_data = fetch.get_all()
        cache.set(key, all_data, timeout=30)

    res = [["created-date", "updated-date", "uuid", "cidr-boot", "mac-boot"]]
    for i in all_data:
        sub_list = list()
        sub_list.append(i["boot-info"]["created-date"])
        sub_list.append(i["boot-info"]["updated-date"])
        sub_list.append(i["boot-info"]["uuid"])
        for j in i["interfaces"]:
            if j["as_boot"]:
                sub_list.append(j["cidrv4"])
                sub_list.append(j["mac"])
        res.append(sub_list)

    return jsonify(res)

if __name__ == "__main__":
    app.logger.setLevel("DEBUG")
    application.run(debug=True)
<|endoftext|>