<commit_before># Copyright Least Authority Enterprises.
# See LICENSE for details.

"""
xUnit TestCase for txkube testing.
"""

from os import environ

from hypothesis import (
    HealthCheck,
    settings,
)

from fixtures import CompoundFixture

from testtools import TestCase as TesttoolsTestCase
from testtools.twistedsupport import AsynchronousDeferredRunTest

from twisted.python.failure import Failure

from ._eliot import CaptureEliotLogs

def _setup_hypothesis():
    settings.register_profile(
        "ci",
        suppress_health_check=[
            # CPU resources available to CI builds typically varies
            # significantly from run to run making it difficult to determine
            # if "too slow" data generation is a result of the code or the
            # execution environment.  Prevent these checks from
            # (intermittently) failing tests that are otherwise fine.
            HealthCheck.too_slow,
        ],
        # By the same reasoning as above, disable the deadline check.
        deadline=None,
    )

    profile_name = environ.get("TXKUBE_HYPOTHESIS_PROFILE", None)
    if profile_name is not None:
        settings.load_profile(profile_name)
_setup_hypothesis()

class AsynchronousDeferredRunTest(AsynchronousDeferredRunTest):
    """
    An asynchronous runner supporting Eliot.
    """
    def _get_log_fixture(self):
        """
        Add ``CaptureEliotLogs`` to the log fixtures which receive special
        treatment so as to be "cleaned up" in the timeout case.

        This ensures eliot logs are reported when tests time out - something
        that will not happen using the normal ``useFixture`` API.

        See <https://bugs.launchpad.net/testtools/+bug/897196>.
        """
        return CompoundFixture([
            super(AsynchronousDeferredRunTest, self)._get_log_fixture(),
            CaptureEliotLogs(),
        ])

class TestCase(TesttoolsTestCase):
    """
    A base class for test cases which automatically uses the
    ``CaptureEliotLogs`` fixture.
    """
    # expectThat and Hypothesis don't communicate well about when the
    # test has failed.  Give them a little help.  These two Hypothesis
    # hooks will check for a flag that testtools sets when it thinks
    # the test has failed and turn it into something Hypothesis can
    # recognize.
    def setup_example(self):
        try:
            # TesttoolsTestCase starts without this attribute set at all.  Get
            # us back to that state.  It won't be set at all on the first
            # setup_example call, nor if the previous run didn't have a failed
            # expectation.
            del self.force_failure
        except AttributeError:
            pass

    def teardown_example(self, ignored):
        if getattr(self, "force_failure", False):
            # If testtools things it's time to stop, translate this into a
            # test failure exception that Hypothesis can see.  This lets
            # Hypothesis know when it has found a falsifying example.  Without
            # it, Hypothesis can't see which of its example runs caused
            # problems.
            self.fail("expectation failed")

def assertNoResult(case, d):
    """
    Assert that ``d`` does not have a result at this point.
    """
    result = []
    d.addBoth(result.append)
    if result:
        if isinstance(result[0], Failure):
            result[0].raiseException()
        else:
            case.fail("Got {} but expected no result".format(result[0]))
<commit_msg>Hide the branch to improve the coverage report<commit_after># Copyright Least Authority Enterprises.
# See LICENSE for details.

"""
xUnit TestCase for txkube testing.
"""

from os import environ

from hypothesis import (
    HealthCheck,
    settings,
)

from fixtures import CompoundFixture

from testtools import TestCase as TesttoolsTestCase
from testtools.twistedsupport import AsynchronousDeferredRunTest

from twisted.python.failure import Failure

from ._eliot import CaptureEliotLogs

def _setup_hypothesis():
    settings.register_profile(
        "ci",
        suppress_health_check=[
            # CPU resources available to CI builds typically varies
            # significantly from run to run making it difficult to determine
            # if "too slow" data generation is a result of the code or the
            # execution environment.  Prevent these checks from
            # (intermittently) failing tests that are otherwise fine.
            HealthCheck.too_slow,
        ],
        # By the same reasoning as above, disable the deadline check.
        deadline=None,
    )
    settings.load_profile(environ.get("TXKUBE_HYPOTHESIS_PROFILE", "default"))
_setup_hypothesis()

class AsynchronousDeferredRunTest(AsynchronousDeferredRunTest):
    """
    An asynchronous runner supporting Eliot.
    """
    def _get_log_fixture(self):
        """
        Add ``CaptureEliotLogs`` to the log fixtures which receive special
        treatment so as to be "cleaned up" in the timeout case.

        This ensures eliot logs are reported when tests time out - something
        that will not happen using the normal ``useFixture`` API.

        See <https://bugs.launchpad.net/testtools/+bug/897196>.
        """
        return CompoundFixture([
            super(AsynchronousDeferredRunTest, self)._get_log_fixture(),
            CaptureEliotLogs(),
        ])

class TestCase(TesttoolsTestCase):
    """
    A base class for test cases which automatically uses the
    ``CaptureEliotLogs`` fixture.
    """
    # expectThat and Hypothesis don't communicate well about when the
    # test has failed.  Give them a little help.  These two Hypothesis
    # hooks will check for a flag that testtools sets when it thinks
    # the test has failed and turn it into something Hypothesis can
    # recognize.
    def setup_example(self):
        try:
            # TesttoolsTestCase starts without this attribute set at all.  Get
            # us back to that state.  It won't be set at all on the first
            # setup_example call, nor if the previous run didn't have a failed
            # expectation.
            del self.force_failure
        except AttributeError:
            pass

    def teardown_example(self, ignored):
        if getattr(self, "force_failure", False):
            # If testtools things it's time to stop, translate this into a
            # test failure exception that Hypothesis can see.  This lets
            # Hypothesis know when it has found a falsifying example.  Without
            # it, Hypothesis can't see which of its example runs caused
            # problems.
            self.fail("expectation failed")

def assertNoResult(case, d):
    """
    Assert that ``d`` does not have a result at this point.
    """
    result = []
    d.addBoth(result.append)
    if result:
        if isinstance(result[0], Failure):
            result[0].raiseException()
        else:
            case.fail("Got {} but expected no result".format(result[0]))
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import shapefile
from shapely.geometry import asShape, mapping
from centerline import Centerline

input_hallways = "../geodata/shp/e01_hallways_union.shp"

def write_geojson(outfilename, indata):
    with open(outfilename, "w") as file_out:
        file_out.write(json.dumps(indata))

def create_shapes(shapefile_path):
    '''
    Create our Polygon
    :param shapefile_path: full path to shapefile
    :return: list of Shapely geometries
    '''
    in_ply = shapefile.Reader(shapefile_path)
    ply_shp = in_ply.shapes()
    print type(ply_shp)

    out_multi_ply = [asShape(feature) for feature in ply_shp]

    print "converting to MultiPolygon: "

    return out_multi_ply

# run our function to create Shapely geometries
shply_ply_halls = create_shapes(input_hallways)

def generate_centerlines(polygon_shps):
    '''
    Create centerlines
    :param polygon_shps: input polygons
    :return: dictionary of linestrings
    '''
    dct_centerlines = {}
    print polygon_shps
    print "i see"

    for i, geom in enumerate(polygon_shps):
        print " now running Centerline creation"
        center_obj = Centerline(geom, 0.5)
        center_line_shply_line = center_obj.create_centerline()
        dct_centerlines[i] = center_line_shply_line

    return dct_centerlines

res_centerlines = generate_centerlines(shply_ply_halls)

def export_center(geojs_file, centerlines):
    '''
    Write output to GeoJSON file
    :param centerlines: input dictionary of linestrings
    :return: write to GeoJSON file
    '''
    with open(geojs_file, 'w') as out:

        for i, key in enumerate(centerlines):
            geom = centerlines[key]
            newline = {'id': key, 'geometry': mapping(geom), 'properties': {'id': key}}

            out.write(json.dumps(newline))

print "now creating centerlines geojson"

outgeojs_file = '../geodata/04_centerline_results_final.geojson'

# write the output GeoJSON file to disk
export_center(outgeojs_file, res_centerlines)

print "FINISHED !  Finally"<commit_msg>renamed input shapefile to smaller hallways polygon<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import shapefile
from shapely.geometry import asShape, mapping
from centerline import Centerline

input_hallways = "../geodata/shp/e01_hallways_small_3857.shp"

def write_geojson(outfilename, indata):
    with open(outfilename, "w") as file_out:
        file_out.write(json.dumps(indata))

def create_shapes(shapefile_path):
    '''
    Create our Polygon
    :param shapefile_path: full path to shapefile
    :return: list of Shapely geometries
    '''
    in_ply = shapefile.Reader(shapefile_path)
    ply_shp = in_ply.shapes()
    print type(ply_shp)

    out_multi_ply = [asShape(feature) for feature in ply_shp]

    print "converting to MultiPolygon: "

    return out_multi_ply

# run our function to create Shapely geometries
shply_ply_halls = create_shapes(input_hallways)

def generate_centerlines(polygon_shps):
    '''
    Create centerlines
    :param polygon_shps: input polygons
    :return: dictionary of linestrings
    '''
    dct_centerlines = {}
    print polygon_shps
    print "i see"

    for i, geom in enumerate(polygon_shps):
        print " now running Centerline creation"
        center_obj = Centerline(geom, 0.5)
        center_line_shply_line = center_obj.create_centerline()
        dct_centerlines[i] = center_line_shply_line

    return dct_centerlines

res_centerlines = generate_centerlines(shply_ply_halls)

def export_center(geojs_file, centerlines):
    '''
    Write output to GeoJSON file
    :param centerlines: input dictionary of linestrings
    :return: write to GeoJSON file
    '''
    with open(geojs_file, 'w') as out:

        for i, key in enumerate(centerlines):
            geom = centerlines[key]
            newline = {'id': key, 'geometry': mapping(geom), 'properties': {'id': key}}

            out.write(json.dumps(newline))

print "now creating centerlines geojson"

outgeojs_file = '../geodata/04_centerline_results_final.geojson'

# write the output GeoJSON file to disk
export_center(outgeojs_file, res_centerlines)

print "FINISHED !  Finally"<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    pygments.styles.neon
    ~~~~~~~~~~~~~~~~~~~~

    Neon Pygments style, modified from the Sublime Text/TextMate
    `Neon Color Scheme <https://github.com/MattDMo/Neon-color-scheme>`__

    tango.py used as template, as it included all the known
    Token types, unlike most (if not all) of the styles included in the
    Pygments distribution. However, the final version of neon.py bears
    no resemblance to tango.

    :copyright: Copyright 2013 by Matt Morrison mattdmo@mattdmo.com
    :license: BSD, see LICENSE for details.
"""

from pygments.style import Style
from pygments.token import Keyword, Name, Comment, String, Error, \
     Number, Operator, Generic, Whitespace, Punctuation, Other, Literal

class NeonStyle(Style):
    """
    Neon style inspired by the Sublime Text Neon Color Scheme -
    https://github.com/MattDMo/Neon-color-scheme
    """

    # work in progress...

    background_color = "#000000"
    highlight_color  = "#2d2d2d"
    default_style    = ""

    styles = {
        # token                    http://pygments.org/docs/styles/#style-rules
        Text:                      "#ffffff", # class:  ''
        Whitespace:                "#f8f8f8",      # class: 'w'
        Error:                     "#f8f8f8 bg:#800f00", # class: 'err'
        Other:                     "#ffffff",                # class 'x'

        Comment:                   "italic #7f817e", # class: 'c'
        Comment.Multiline:         "italic #7f817e", # class: 'cm'
        Comment.Preproc:           "italic #7f817e", # class: 'cp'
        Comment.Single:            "italic #7f817e", # class: 'c1'
        Comment.Special:           "italic #7f817e", # class: 'cs'

        Keyword:                   "#0b93ff",   # class: 'k'
        Keyword.Constant:          "italic #ff07a0",   # class: 'kc'
        Keyword.Declaration:       "#0b93ff",   # class: 'kd'
        Keyword.Namespace:         "#0b93ff",   # class: 'kn'
        Keyword.Pseudo:            "#0b93ff",   # class: 'kp'
        Keyword.Reserved:          "#0b93ff",   # class: 'kr'
        Keyword.Type:              "#0b93ff",   # class: 'kt'

        Operator:                  "bold #a7a3ff",   # class: 'o'
        Operator.Word:             "bold #a7a3ff",   # class: 'ow' - like keywords

        Punctuation:               "#fffef7",   # class: 'p'

        Name:                      "#08d879",        # class: 'n'
        Name.Attribute:            "#c4a000",        # class: 'na' - to be revised
        Name.Builtin:              "italic #e0a1ff",        # class: 'nb'
        Name.Builtin.Pseudo:       "italic #0aedff",        # class: 'bp'
        Name.Class:                "#cfff01",        # class: 'nc' - to be revised
        Name.Constant:             "italic #eb939a",        # class: 'no' - to be revised
        Name.Decorator:            "#b6b8fe bg:#070e48",   # class: 'nd' - to be revised
        Name.Entity:               "#ce5c00",        # class: 'ni'
        Name.Exception:            "bold #ff1e00",   # class: 'ne'
        Name.Function:             "#10ff02",        # class: 'nf'
        Name.Property:             "#b6b8fe",        # class: 'py'
        Name.Label:                "#f57900",        # class: 'nl'
        Name.Namespace:            "#11bd7c",        # class: 'nn' - to be revised
        Name.Other:                "#ff25d9",        # class: 'nx'
        Name.Tag:                  "bold #02aeff",   # class: 'nt' - like a keyword
        Name.Variable:             "#ff25d9",        # class: 'nv' - to be revised
        Name.Variable.Class:       "#cfff01",        # class: 'vc' - to be revised
        Name.Variable.Global:      "#d285cc",        # class: 'vg' - to be revised
        Name.Variable.Instance:    "#d285cc",        # class: 'vi' - to be revised

        Number:                    "#ff0604",   # class: 'm'
        Number.Float:              "#ff0604",   # class: 'mf'
        Number.Hex:                "italic #ff0604",   # class: 'mh'
        Number.Integer:            "#ff0604",   # class: 'mi'
        Number.Integer.Long:       "#ff0604",   # class: 'il'
        Number.Oct:                "italic #ff0604",   # class: 'mo'

        Literal:                   "#db1e44",        # class: 'l'
        Literal.Date:              "#14ff01",        # class: 'ld'

        String:                    "#ffdf02",        # class: 's'
        String.Backtick:           "italic #ffdf02",        # class: 'sb'
        String.Char:               "#ffdf02",        # class: 'sc'
        String.Doc:                "italic #218b97", # class: 'sd' - like a comment
        String.Double:             "#ffdf02",        # class: 's2'
        String.Escape:             "italic #ff087b",        # class: 'se'
        String.Heredoc:            "italic #3a771f",        # class: 'sh'
        String.Interpol:           "#ffdf02",        # class: 'si'
        String.Other:              "#ffdf02",        # class: 'sx'
        String.Regex:              "#ffe4a6",        # class: 'sr'
        String.Single:             "#ffdf02",        # class: 's1'
        String.Symbol:             "#ff9705",        # class: 'ss'

        Generic:                   "#ffffff",        # class: 'g'
        Generic.Deleted:           "#f8f8f8 bg:#4c0900",        # class: 'gd'
        Generic.Emph:              "bold italic #ffffff", # class: 'ge'
        Generic.Error:             "bold #ff1e00",        # class: 'gr'
        Generic.Heading:           "bold #000080",   # class: 'gh'
        Generic.Inserted:          "#f8f8f8 bg:#154f00",        # class: 'gi'
        Generic.Output:            "italic #ffffff", # class: 'go'
        Generic.Prompt:            "#22ff31",        # class: 'gp'
        Generic.Strong:            "bold #ffffff",   # class: 'gs'
        Generic.Subheading:        "bold #ff00ff",   # class: 'gu'
        Generic.Traceback:         "bold #f00bad",   # class: 'gt'
    }
<commit_msg>apparently Text is not defined...<commit_after># -*- coding: utf-8 -*-
"""
    pygments.styles.neon
    ~~~~~~~~~~~~~~~~~~~~

    Neon Pygments style, modified from the Sublime Text/TextMate
    `Neon Color Scheme <https://github.com/MattDMo/Neon-color-scheme>`__

    tango.py used as template, as it included all the known
    Token types, unlike most (if not all) of the styles included in the
    Pygments distribution. However, the final version of neon.py bears
    no resemblance to tango.

    :copyright: Copyright 2013 by Matt Morrison mattdmo@mattdmo.com
    :license: BSD, see LICENSE for details.
"""

from pygments.style import Style
from pygments.token import Keyword, Name, Comment, String, Error, \
     Number, Operator, Generic, Whitespace, Punctuation, Other, Literal

class NeonStyle(Style):
    """
    Neon style inspired by the Sublime Text Neon Color Scheme -
    https://github.com/MattDMo/Neon-color-scheme
    """

    # work in progress...

    background_color = "#000000"
    highlight_color  = "#2d2d2d"
    default_style    = ""

    styles = {
        # token                    http://pygments.org/docs/styles/#style-rules
        # Text:                      "#ffffff", # class:  ''
        Whitespace:                "#f8f8f8",      # class: 'w'
        Error:                     "#f8f8f8 bg:#800f00", # class: 'err'
        Other:                     "#ffffff",                # class 'x'

        Comment:                   "italic #7f817e", # class: 'c'
        Comment.Multiline:         "italic #7f817e", # class: 'cm'
        Comment.Preproc:           "italic #7f817e", # class: 'cp'
        Comment.Single:            "italic #7f817e", # class: 'c1'
        Comment.Special:           "italic #7f817e", # class: 'cs'

        Keyword:                   "#0b93ff",   # class: 'k'
        Keyword.Constant:          "italic #ff07a0",   # class: 'kc'
        Keyword.Declaration:       "#0b93ff",   # class: 'kd'
        Keyword.Namespace:         "#0b93ff",   # class: 'kn'
        Keyword.Pseudo:            "#0b93ff",   # class: 'kp'
        Keyword.Reserved:          "#0b93ff",   # class: 'kr'
        Keyword.Type:              "#0b93ff",   # class: 'kt'

        Operator:                  "bold #a7a3ff",   # class: 'o'
        Operator.Word:             "bold #a7a3ff",   # class: 'ow' - like keywords

        Punctuation:               "#fffef7",   # class: 'p'

        Name:                      "#08d879",        # class: 'n'
        Name.Attribute:            "#c4a000",        # class: 'na' - to be revised
        Name.Builtin:              "italic #e0a1ff",        # class: 'nb'
        Name.Builtin.Pseudo:       "italic #0aedff",        # class: 'bp'
        Name.Class:                "#cfff01",        # class: 'nc' - to be revised
        Name.Constant:             "italic #eb939a",        # class: 'no' - to be revised
        Name.Decorator:            "#b6b8fe bg:#070e48",   # class: 'nd' - to be revised
        Name.Entity:               "#ce5c00",        # class: 'ni'
        Name.Exception:            "bold #ff1e00",   # class: 'ne'
        Name.Function:             "#10ff02",        # class: 'nf'
        Name.Property:             "#b6b8fe",        # class: 'py'
        Name.Label:                "#f57900",        # class: 'nl'
        Name.Namespace:            "#11bd7c",        # class: 'nn' - to be revised
        Name.Other:                "#ff25d9",        # class: 'nx'
        Name.Tag:                  "bold #02aeff",   # class: 'nt' - like a keyword
        Name.Variable:             "#ff25d9",        # class: 'nv' - to be revised
        Name.Variable.Class:       "#cfff01",        # class: 'vc' - to be revised
        Name.Variable.Global:      "#d285cc",        # class: 'vg' - to be revised
        Name.Variable.Instance:    "#d285cc",        # class: 'vi' - to be revised

        Number:                    "#ff0604",   # class: 'm'
        Number.Float:              "#ff0604",   # class: 'mf'
        Number.Hex:                "italic #ff0604",   # class: 'mh'
        Number.Integer:            "#ff0604",   # class: 'mi'
        Number.Integer.Long:       "#ff0604",   # class: 'il'
        Number.Oct:                "italic #ff0604",   # class: 'mo'

        Literal:                   "#db1e44",        # class: 'l'
        Literal.Date:              "#14ff01",        # class: 'ld'

        String:                    "#ffdf02",        # class: 's'
        String.Backtick:           "italic #ffdf02",        # class: 'sb'
        String.Char:               "#ffdf02",        # class: 'sc'
        String.Doc:                "italic #218b97", # class: 'sd' - like a comment
        String.Double:             "#ffdf02",        # class: 's2'
        String.Escape:             "italic #ff087b",        # class: 'se'
        String.Heredoc:            "italic #3a771f",        # class: 'sh'
        String.Interpol:           "#ffdf02",        # class: 'si'
        String.Other:              "#ffdf02",        # class: 'sx'
        String.Regex:              "#ffe4a6",        # class: 'sr'
        String.Single:             "#ffdf02",        # class: 's1'
        String.Symbol:             "#ff9705",        # class: 'ss'

        Generic:                   "#ffffff",        # class: 'g'
        Generic.Deleted:           "#f8f8f8 bg:#4c0900",        # class: 'gd'
        Generic.Emph:              "bold italic #ffffff", # class: 'ge'
        Generic.Error:             "bold #ff1e00",        # class: 'gr'
        Generic.Heading:           "bold #000080",   # class: 'gh'
        Generic.Inserted:          "#f8f8f8 bg:#154f00",        # class: 'gi'
        Generic.Output:            "italic #ffffff", # class: 'go'
        Generic.Prompt:            "#22ff31",        # class: 'gp'
        Generic.Strong:            "bold #ffffff",   # class: 'gs'
        Generic.Subheading:        "bold #ff00ff",   # class: 'gu'
        Generic.Traceback:         "bold #f00bad",   # class: 'gt'
    }
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

import sys
import json
import argparse
import tempfile
import shutil
import os

import random
import hashlib
import datetime
import docker
import logging
import tarfile
import cStringIO

# Module level docker connection
# TODO: This could be made configurable later
DOCKER_CLIENT = docker.Client(
    base_url='unix://var/run/docker.sock',
    timeout=240)

# Module level logger
LOG = logging.getLogger()
_HANDLER = logging.StreamHandler()
_FORMATTER = logging.Formatter('%(asctime)s %(name)-12s %(levelname)-8s %(message)s')
_HANDLER.setFormatter(_FORMATTER)
LOG.addHandler(_HANDLER)
LOG.setLevel(logging.DEBUG)

class Chdir:
  """ Context manager for changing the current working directory """
  def __init__(self, newPath):
    self.newPath = os.path.expanduser(newPath)

  def __enter__(self):
    self.savedPath = os.getcwd()
    os.chdir(self.newPath)

  def __exit__(self, etype, value, traceback):
    os.chdir(self.savedPath)

def _read_layers(layers, image_id):
  """ Reads the JSON metadata for specified layer / image id """

  layer = DOCKER_CLIENT.inspect_image(image_id)
  layers.append(layer['Id'])

  if 'Parent' in layer and layer['Parent']:
    _read_layers(layers, layer['Parent'])

def _save_image(image_id, tar_file):
  """ Saves the image as a tar archive under specified name """

  LOG.debug("Saving image %s to %s file..." % (image_id, tar_file))

  image = DOCKER_CLIENT.get_image(image_id)

  with open(tar_file, 'w') as f:
    f.write(image.data)

  LOG.debug("Image saved!")

def _unpack(tar_file, directory):
  """ Unpacks tar archive to selected directory """

  LOG.debug("Unpacking %s tar file to %s directory" % (tar_file, directory))

  with tarfile.open(tar_file, 'r') as tar:
    tar.extractall(path=directory)

  LOG.debug("Archive unpacked!")

def _move_unmodified_layers(layers, squash_id, src, dest):
  """
  This moves all the layers that should be copied as-is.
  In other words - all layers that are not meant to be squashed will be
  moved from the old image to the new image untouched.
  """
  for layer in layers:
    LOG.debug("Moving unmodified layer %s..." % layer)
    shutil.move(os.path.join(src, layer), dest)
    if layer == squash_id:
      # Stop if we are at the first layer that was squashed
      return

def _marker_files(tar, layer_id):
  """
  Searches for marker files in the specified archive.

  Docker marker files are files taht have the .wh. prefix in the name.
  These files mark the corresponding file to be removed (hidden) when
  we start a container from the image.
  """
  marker_files = {}

  LOG.debug("Searching for marker files in '%s' archive..." % tar.name)

  for member in tar.getmembers():
    if '.wh.' in member.name:
      LOG.debug("Found '%s' marker file" % member.name)
      marker_files[member.name] = member

  if marker_files:
    LOG.debug("Following files are marked to skip in the %s layer: %s" % (layer_id, " ".join(marker_files.keys())))

  return marker_files

def _generate_target_json(old_image_id, new_image_id, squash_id, squashed_dir):
  json_file = os.path.join(squashed_dir, "json")
  squashed_tar = os.path.join(squashed_dir, "layer.tar")
  # Read the original metadata
  metadata = DOCKER_CLIENT.inspect_image(old_image_id)

  # Update the fields
  metadata['Id'] = new_image_id
  metadata['Parent'] = squash_id
  metadata['Config']['Image'] = squash_id
  metadata['Created'] = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S.%fZ')
  metadata['Size'] = os.path.getsize(squashed_tar)

  # Remove unnecessary fields
  del metadata['ContainerConfig']
  del metadata['Container']
  del metadata['Config']['Hostname']

  with open(json_file, 'w') as f:
    json.dump(metadata, f)

def _generate_repositories_json(repositories_file, new_image_id, tag):
  if ':' in tag:
    name, tag = tag.split(':')
  else:
    name = tag
    tag = "latest"

  repos = {}
  repos[name] = {}
  repos[name][tag] = new_image_id

  with open(repositories_file, 'w') as f:
    json.dump(repos, f)

def _load_image(directory):
  c = cStringIO.StringIO()

  with tarfile.open(mode='w', fileobj=c) as tar:
    LOG.debug("Generating tar archive for the squashed image...")
    with Chdir(directory):
      tar.add(".")
    LOG.debug("Archive generated")

  LOG.debug("Uploading image...")
  DOCKER_CLIENT.load_image(c.getvalue())
  LOG.debug("Image uploaded")

  c.close()

def _layers_to_squash(layers, from_layer):
  """ Prepares a list of layer IDs that should be squashed """
  to_squash = []

  for l in reversed(layers):
    if l == from_layer:
      break

    to_squash.append(l)

  to_squash.reverse()

  return to_squash

def _prepare_tmp_directory(provided_tmp_dir):
  """ Creates temporary directory that is used to work on layers """
  if provided_tmp_dir:
    if os.path.exists(provided_tmp_dir):
      LOG.error("The '%s' directory already exists, please remove it before you proceed, aborting." % provided_tmp_dir)
      sys.exit(1)
    os.makedirs(provided_tmp_dir)
    return provided_tmp_dir
  else:
    return tempfile.mkdtemp(prefix="tmp-docker-squash-")

def _squash_layers(layers_to_squash, squashed_tar_file, old_image_dir):
  # Reverse the layers to squash - we begin with the newest one
  # to make the tar lighter
  layers_to_squash.reverse()

  LOG.info("Starting squashing...")

  with tarfile.open(squashed_tar_file, 'w') as squashed_tar:
    unskipped_markers = {}

    for layer_id in layers_to_squash:
      layer_tar_file = os.path.join(old_image_dir, layer_id, "layer.tar")

      LOG.info("Squashing layer %s..." % layer_id)

      # Open the exiting layer to squash
      with tarfile.open(layer_tar_file, 'r') as layer_tar:
        # Find all marker files for all layers
        markers = _marker_files(layer_tar, layer_id)
        tar_files = [o.name for o in layer_tar.getmembers()]

        to_skip = []

        for marker_name in unskipped_markers.keys():
          actual_file = marker_name.replace('.wh.', '')

          if actual_file in tar_files:
            to_skip.append(marker_name)
            to_skip.append(actual_file)
            del(unskipped_markers[marker_name])

        for marker_name, marker in markers.iteritems():
          actual_file = marker_name.replace('.wh.', '')
          to_skip.append(marker_name)

          if actual_file in tar_files:
            to_skip.append(actual_file)
          else:
            # We can safely add the file content, because marker files are empty
            unskipped_markers[marker_name] = { 'file': layer_tar.extractfile(member), 'info': marker }

        if to_skip:
          LOG.debug("Following files are marked to skip when squashing layer %s: %s" % (layer_id, to_skip))

        # Copy all the files to the new tar
        for member in layer_tar.getmembers():
          # Skip files that are marked to be skipped
          if member.name in to_skip:
            LOG.debug("Skipping '%s' file because it's on the list to skip files" % member.name)
            continue

          # List of filenames in the squashed archive
          squashed_files = [o.name for o in squashed_tar.getmembers()]

          # Check if file is already added to the archive
          if member.name in squashed_files:
            # File already exist in the squashed archive, skip it because
            # file want to add is older than the one already in the archive.
            # This is true because we do reverse squashing - from newer to older layer
            LOG.debug("Skipping '%s' file because it's older than file already added to the archive" % member.name)
            continue

          if member.issym():
            # Special case: symlinks
            squashed_tar.addfile(member)
          else:
            # Finally add the file to archive
            squashed_tar.addfile(member, layer_tar.extractfile(member))

    # We copied all the files from all layers, but if there are
    # still some marker files - we need to add them back because these
    # remove (technically: hide) files from layers unaffected
    # by squashing
    for marker_name, marker in unskipped_markers.iteritems():
      LOG.debug("Adding '%s' marker file back since the file it refers to was not found in any layers we squashed..." % marker_name)
      squashed_tar.addfile(marker['info'], marker['file'])

  LOG.debug("Squashing done!")

def main(args):

  # The image id or name of the image to be squashed
  try:
    old_image_id = DOCKER_CLIENT.inspect_image(args.image)['Id']
  except:
    LOG.error("Could not get the image ID to squash, please check provided 'image' argument: %s" % args.image)
    sys.exit(1)

  # The id or name of the layer/image that the squashing should begin from
  # This layer WILL NOT be squashed, but all next layers will
  try:
    squash_id = DOCKER_CLIENT.inspect_image(args.layer)['Id']
  except:
    LOG.error("Could not get the layer ID to squash, please check provided 'layer' argument: %s" % args.layer)
    sys.exit(1)

  old_layers = []

  # Read all layers in the image
  _read_layers(old_layers, old_image_id)

  old_layers.reverse()

  LOG.info("Old image has %s layers", len(old_layers))
  LOG.debug("Old layers: %s", old_layers)

  if not squash_id in old_layers:
    LOG.error("Couldn't find the provided layer (%s) in the %s image" % (args.layer, args.image))
    sys.exit(1)

  # Find the layers to squash
  layers_to_squash = _layers_to_squash(old_layers, squash_id)

  LOG.info("We have %s layers to squash", len(layers_to_squash))
  LOG.debug("Layers to squash: %s", layers_to_squash)

  if len(layers_to_squash) == 0:
    LOG.error("There are no layers to squash, aborting")
    sys.exit(1)

  # Prepare temporary directory where all the work will be executed
  tmp_dir = _prepare_tmp_directory(args.tmp_dir)

  # Location of the tar with the old image
  old_image_tar = os.path.join(tmp_dir, "image.tar")

  # Save the image in tar format in the tepmorary directory
  _save_image(old_image_id, old_image_tar)

  # Directory where the old layers will be unpacked
  old_image_dir = os.path.join(tmp_dir, "old")
  os.makedirs(old_image_dir)

  # Unpack the image
  LOG.info("Unpacking exported tar (%s)..." % old_image_tar)
  _unpack(old_image_tar, old_image_dir)

  # Remove the tar file early to save some space
  LOG.info("Removing exported tar (%s)..." % old_image_tar)
  os.remove(old_image_tar)

  # Directory where the new layers will be unpacked in prepareation to
  # import it to Docker
  new_image_dir = os.path.join(tmp_dir, "new")
  os.makedirs(new_image_dir)

  # Generate a new image id for the squashed layer
  new_image_id = hashlib.sha256(str(random.getrandbits(128))).hexdigest()

  LOG.info("New layer ID for squashed content will be: %s" % new_image_id)

  # Prepare a directory for squashed layer content
  squashed_dir = os.path.join(new_image_dir, new_image_id)
  os.makedirs(squashed_dir)

  # Location of the tar archive with the squashed layers
  squashed_tar = os.path.join(squashed_dir, "layer.tar")

  # Append all the layers on each other
  _squash_layers(layers_to_squash, squashed_tar, old_image_dir)

  # Move all the layers that should be untouched
  _move_unmodified_layers(old_layers, squash_id, old_image_dir, new_image_dir)

  # Generate the metadata JSON based on the original one
  _generate_target_json(old_image_id, new_image_id, squash_id, squashed_dir)

  # Generate the metadata JSON with information about the images
  _generate_repositories_json(os.path.join(new_image_dir, "repositories"), new_image_id, args.tag)

  # And finally tar everything up and load into Docker
  _load_image(new_image_dir)

  # Cleanup the temporary directory
  shutil.rmtree(tmp_dir)

if __name__ == "__main__":
  PARSER = argparse.ArgumentParser(description='Squashes all layers in the image from the layer specified as "layer" argument.')
  PARSER.add_argument('image', help='Image to be squashed')
  PARSER.add_argument('layer', help='ID of the layer or image ID or image name')
  PARSER.add_argument('tag', help='Specify the tag to be used for the new image')
  PARSER.add_argument('-t', '--tmp-dir', help='Temporary directory to be used')
  ARGS = PARSER.parse_args()

  main(ARGS)
<commit_msg>Friendlier logging, use -v for more logs<commit_after># -*- coding: utf-8 -*-

import sys
import json
import argparse
import tempfile
import shutil
import os

import random
import hashlib
import datetime
import docker
import logging
import tarfile
import cStringIO

# Module level docker connection
# TODO: This could be made configurable later
DOCKER_CLIENT = docker.Client(
    base_url='unix://var/run/docker.sock',
    timeout=240)

# Module level logger
LOG = logging.getLogger()
_HANDLER = logging.StreamHandler()
_FORMATTER = logging.Formatter('%(asctime)s %(name)-12s %(levelname)-8s %(message)s')
_HANDLER.setFormatter(_FORMATTER)
LOG.addHandler(_HANDLER)

class Chdir:
  """ Context manager for changing the current working directory """
  def __init__(self, newPath):
    self.newPath = os.path.expanduser(newPath)

  def __enter__(self):
    self.savedPath = os.getcwd()
    os.chdir(self.newPath)

  def __exit__(self, etype, value, traceback):
    os.chdir(self.savedPath)

def _read_layers(layers, image_id):
  """ Reads the JSON metadata for specified layer / image id """

  layer = DOCKER_CLIENT.inspect_image(image_id)
  layers.append(layer['Id'])

  if 'Parent' in layer and layer['Parent']:
    _read_layers(layers, layer['Parent'])

def _save_image(image_id, tar_file):
  """ Saves the image as a tar archive under specified name """

  LOG.info("Saving image %s to %s file..." % (image_id, tar_file))

  image = DOCKER_CLIENT.get_image(image_id)

  with open(tar_file, 'w') as f:
    f.write(image.data)

  LOG.info("Image saved!")

def _unpack(tar_file, directory):
  """ Unpacks tar archive to selected directory """

  LOG.info("Unpacking %s tar file to %s directory" % (tar_file, directory))

  with tarfile.open(tar_file, 'r') as tar:
    tar.extractall(path=directory)

  # Remove the tar file early to save some space
  LOG.debug("Removing exported tar (%s)..." % tar_file)
  os.remove(tar_file)

  LOG.info("Archive unpacked!")

def _move_unmodified_layers(layers, squash_id, src, dest):
  """
  This moves all the layers that should be copied as-is.
  In other words - all layers that are not meant to be squashed will be
  moved from the old image to the new image untouched.
  """
  for layer in layers:
    LOG.debug("Moving unmodified layer %s..." % layer)
    shutil.move(os.path.join(src, layer), dest)
    if layer == squash_id:
      # Stop if we are at the first layer that was squashed
      return

def _marker_files(tar, layer_id):
  """
  Searches for marker files in the specified archive.

  Docker marker files are files taht have the .wh. prefix in the name.
  These files mark the corresponding file to be removed (hidden) when
  we start a container from the image.
  """
  marker_files = {}

  LOG.debug("Searching for marker files in '%s' archive..." % tar.name)

  for member in tar.getmembers():
    if '.wh.' in member.name:
      LOG.debug("Found '%s' marker file" % member.name)
      marker_files[member.name] = member

  if marker_files:
    LOG.debug("Following files are marked to skip in the %s layer: %s" % (layer_id, " ".join(marker_files.keys())))

  return marker_files

def _generate_target_json(old_image_id, new_image_id, squash_id, squashed_dir):
  json_file = os.path.join(squashed_dir, "json")
  squashed_tar = os.path.join(squashed_dir, "layer.tar")
  # Read the original metadata
  metadata = DOCKER_CLIENT.inspect_image(old_image_id)

  # Update the fields
  metadata['Id'] = new_image_id
  metadata['Parent'] = squash_id
  metadata['Config']['Image'] = squash_id
  metadata['Created'] = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S.%fZ')
  metadata['Size'] = os.path.getsize(squashed_tar)

  # Remove unnecessary fields
  del metadata['ContainerConfig']
  del metadata['Container']
  del metadata['Config']['Hostname']

  with open(json_file, 'w') as f:
    json.dump(metadata, f)

def _generate_repositories_json(repositories_file, new_image_id, tag):
  if ':' in tag:
    name, tag = tag.split(':')
  else:
    name = tag
    tag = "latest"

  repos = {}
  repos[name] = {}
  repos[name][tag] = new_image_id

  with open(repositories_file, 'w') as f:
    json.dump(repos, f)

def _load_image(directory):
  c = cStringIO.StringIO()

  with tarfile.open(mode='w', fileobj=c) as tar:
    LOG.debug("Generating tar archive for the squashed image...")
    with Chdir(directory):
      tar.add(".")
    LOG.debug("Archive generated")

  LOG.info("Loading squashed image...")
  DOCKER_CLIENT.load_image(c.getvalue())
  LOG.info("Image loaded!")

  c.close()

def _layers_to_squash(layers, from_layer):
  """ Prepares a list of layer IDs that should be squashed """
  to_squash = []

  for l in reversed(layers):
    if l == from_layer:
      break

    to_squash.append(l)

  to_squash.reverse()

  return to_squash

def _prepare_tmp_directory(provided_tmp_dir):
  """ Creates temporary directory that is used to work on layers """
  if provided_tmp_dir:
    if os.path.exists(provided_tmp_dir):
      LOG.error("The '%s' directory already exists, please remove it before you proceed, aborting." % provided_tmp_dir)
      sys.exit(1)
    os.makedirs(provided_tmp_dir)
    return provided_tmp_dir
  else:
    return tempfile.mkdtemp(prefix="tmp-docker-squash-")

def _squash_layers(layers_to_squash, squashed_tar_file, old_image_dir):
  # Reverse the layers to squash - we begin with the newest one
  # to make the tar lighter
  layers_to_squash.reverse()

  LOG.info("Starting squashing...")

  with tarfile.open(squashed_tar_file, 'w') as squashed_tar:
    unskipped_markers = {}

    for layer_id in layers_to_squash:
      layer_tar_file = os.path.join(old_image_dir, layer_id, "layer.tar")

      LOG.info("Squashing layer %s..." % layer_id)

      # Open the exiting layer to squash
      with tarfile.open(layer_tar_file, 'r') as layer_tar:
        # Find all marker files for all layers
        markers = _marker_files(layer_tar, layer_id)
        tar_files = [o.name for o in layer_tar.getmembers()]

        to_skip = []

        for marker_name in unskipped_markers.keys():
          actual_file = marker_name.replace('.wh.', '')

          if actual_file in tar_files:
            to_skip.append(marker_name)
            to_skip.append(actual_file)
            del(unskipped_markers[marker_name])

        for marker_name, marker in markers.iteritems():
          actual_file = marker_name.replace('.wh.', '')
          to_skip.append(marker_name)

          if actual_file in tar_files:
            to_skip.append(actual_file)
          else:
            # We can safely add the file content, because marker files are empty
            unskipped_markers[marker_name] = { 'file': layer_tar.extractfile(member), 'info': marker }

        if to_skip:
          LOG.debug("Following files are marked to skip when squashing layer %s: %s" % (layer_id, to_skip))

        # Copy all the files to the new tar
        for member in layer_tar.getmembers():
          # Skip files that are marked to be skipped
          if member.name in to_skip:
            LOG.debug("Skipping '%s' file because it's on the list to skip files" % member.name)
            continue

          # List of filenames in the squashed archive
          squashed_files = [o.name for o in squashed_tar.getmembers()]

          # Check if file is already added to the archive
          if member.name in squashed_files:
            # File already exist in the squashed archive, skip it because
            # file want to add is older than the one already in the archive.
            # This is true because we do reverse squashing - from newer to older layer
            LOG.debug("Skipping '%s' file because it's older than file already added to the archive" % member.name)
            continue

          if member.issym():
            # Special case: symlinks
            squashed_tar.addfile(member)
          else:
            # Finally add the file to archive
            squashed_tar.addfile(member, layer_tar.extractfile(member))

    # We copied all the files from all layers, but if there are
    # still some marker files - we need to add them back because these
    # remove (technically: hide) files from layers unaffected
    # by squashing
    for marker_name, marker in unskipped_markers.iteritems():
      LOG.debug("Adding '%s' marker file back since the file it refers to was not found in any layers we squashed..." % marker_name)
      squashed_tar.addfile(marker['info'], marker['file'])

  LOG.debug("Squashing done!")

def main(args):

  if args.verbose:
    LOG.setLevel(logging.DEBUG)
  else:
    LOG.setLevel(logging.INFO)

  # The image id or name of the image to be squashed
  try:
    old_image_id = DOCKER_CLIENT.inspect_image(args.image)['Id']
  except:
    LOG.error("Could not get the image ID to squash, please check provided 'image' argument: %s" % args.image)
    sys.exit(1)

  # The id or name of the layer/image that the squashing should begin from
  # This layer WILL NOT be squashed, but all next layers will
  try:
    squash_id = DOCKER_CLIENT.inspect_image(args.layer)['Id']
  except:
    LOG.error("Could not get the layer ID to squash, please check provided 'layer' argument: %s" % args.layer)
    sys.exit(1)

  if ':' in args.tag:
    tag = args.tag
  else:
    tag = "%s:latest" % args.tag

  LOG.info("Attempting to squash image %s...", old_image_id)

  old_layers = []

  # Read all layers in the image
  _read_layers(old_layers, old_image_id)

  old_layers.reverse()

  LOG.info("Old image has %s layers", len(old_layers))
  LOG.debug("Old layers: %s", old_layers)

  if not squash_id in old_layers:
    LOG.error("Couldn't find the provided layer (%s) in the %s image" % (args.layer, args.image))
    sys.exit(1)

  LOG.info("We will squash from layer %s", squash_id)

  # Find the layers to squash
  layers_to_squash = _layers_to_squash(old_layers, squash_id)

  LOG.info("We have %s layers to squash", len(layers_to_squash))
  LOG.debug("Layers to squash: %s", layers_to_squash)

  if len(layers_to_squash) == 0:
    LOG.error("There are no layers to squash, aborting")
    sys.exit(1)

  # Prepare temporary directory where all the work will be executed
  tmp_dir = _prepare_tmp_directory(args.tmp_dir)

  # Location of the tar with the old image
  old_image_tar = os.path.join(tmp_dir, "image.tar")

  # Save the image in tar format in the tepmorary directory
  _save_image(old_image_id, old_image_tar)

  # Directory where the old layers will be unpacked
  old_image_dir = os.path.join(tmp_dir, "old")
  os.makedirs(old_image_dir)

  # Unpack the image
  _unpack(old_image_tar, old_image_dir)

  # Directory where the new layers will be unpacked in prepareation to
  # import it to Docker
  new_image_dir = os.path.join(tmp_dir, "new")
  os.makedirs(new_image_dir)

  # Generate a new image id for the squashed layer
  new_image_id = hashlib.sha256(str(random.getrandbits(128))).hexdigest()

  LOG.info("New layer ID for squashed content will be: %s" % new_image_id)

  # Prepare a directory for squashed layer content
  squashed_dir = os.path.join(new_image_dir, new_image_id)
  os.makedirs(squashed_dir)

  # Location of the tar archive with the squashed layers
  squashed_tar = os.path.join(squashed_dir, "layer.tar")

  # Append all the layers on each other
  _squash_layers(layers_to_squash, squashed_tar, old_image_dir)

  # Move all the layers that should be untouched
  _move_unmodified_layers(old_layers, squash_id, old_image_dir, new_image_dir)

  # Generate the metadata JSON based on the original one
  _generate_target_json(old_image_id, new_image_id, squash_id, squashed_dir)

  # Generate the metadata JSON with information about the images
  _generate_repositories_json(os.path.join(new_image_dir, "repositories"), new_image_id, tag)

  # And finally tar everything up and load into Docker
  _load_image(new_image_dir)

  # Cleanup the temporary directory
  shutil.rmtree(tmp_dir)

  LOG.info("Finished, image registered as %s", tag)

if __name__ == "__main__":
  PARSER = argparse.ArgumentParser(description='Squashes all layers in the image from the layer specified as "layer" argument.')
  PARSER.add_argument('image', help='Image to be squashed')
  PARSER.add_argument('layer', help='ID of the layer or image ID or image name')
  PARSER.add_argument('tag', help='Specify the tag to be used for the new image')
  PARSER.add_argument('-t', '--tmp-dir', help='Temporary directory to be used')
  PARSER.add_argument('-v', '--verbose', action='store_true', help='Verbose output')
  ARGS = PARSER.parse_args()

  main(ARGS)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import sys
import os
import tempfile
import subprocess
import py
import execnet
from execnet import gateway_base as serializer

def _find_version(suffix=""):
    name = "python" + suffix
    executable = py.path.local.sysfind(name)
    if executable is None:
        if sys.platform == "win32" and suffix == "3":
            for name in ('python31', 'python30'):
                executable = py.path.local(r"c:\\%s\python.exe" % (name,))
                if executable.check():
                    return executable
        py.test.skip("can't find a %r executable" % (name,))
    return executable

def setup_module(mod):
    mod.TEMPDIR = py.path.local(tempfile.mkdtemp())
    if sys.version_info > (3, 0):
        mod._py3_wrapper = PythonWrapper(py.path.local(sys.executable))
        mod._py2_wrapper = PythonWrapper(_find_version())
    else:
        mod._py3_wrapper = PythonWrapper(_find_version("3"))
        mod._py2_wrapper = PythonWrapper(py.path.local(sys.executable))

def teardown_module(mod):
    TEMPDIR.remove(True)

pyimportdir = str(py.path.local(execnet.__file__).dirpath().dirpath())
class PythonWrapper(object):

    def __init__(self, executable):
        self.executable = executable

    def dump(self, obj_rep):
        script_file = TEMPDIR.join("dump.py")
        script_file.write("""
import sys
sys.path.insert(0, %r)
from execnet import gateway_base as serializer
if sys.version_info > (3, 0): # Need binary output
    sys.stdout = sys.stdout.detach()
saver = serializer.Serializer(sys.stdout)
saver.save(%s)""" % (pyimportdir, obj_rep,))
        popen = subprocess.Popen([str(self.executable), str(script_file)],
                                 stdin=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 stdout=subprocess.PIPE)
        ret = popen.wait()
        if ret:
            raise py.process.cmdexec.Error(ret, ret, str(self.executable),
                                           stdout, stderr)
        return popen.stdout.read()

    def load(self, data, option_args="__class__"):
        script_file = TEMPDIR.join("load.py")
        script_file.write(r"""
import sys
sys.path.insert(0, %r)
from execnet import gateway_base as serializer
if sys.version_info > (3, 0):
    sys.stdin = sys.stdin.detach()
loader = serializer.Unserializer(sys.stdin)
loader.%s
obj = loader.load()
sys.stdout.write(type(obj).__name__ + "\n")
sys.stdout.write(repr(obj))""" % (pyimportdir, option_args,))
        popen = subprocess.Popen([str(self.executable), str(script_file)],
                                 stdin=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 stdout=subprocess.PIPE)
        stdout, stderr = popen.communicate(data)
        ret = popen.returncode
        if ret:
            raise py.process.cmdexec.Error(ret, ret, str(self.executable),
                                           stdout, stderr)
        return [s.decode("ascii") for s in stdout.splitlines()]

    def __repr__(self):
        return "<PythonWrapper for %s>" % (self.executable,)

def pytest_funcarg__py2(request):
    return _py2_wrapper

def pytest_funcarg__py3(request):
    return _py3_wrapper

def pytest_funcarg__dump(request):
    py_dump = request.getfuncargvalue(request.param[0])
    return py_dump.dump

def pytest_funcarg__load(request):
    py_dump = request.getfuncargvalue(request.param[1])
    return py_dump.load

def pytest_generate_tests(metafunc):
    if 'dump' in metafunc.funcargnames and 'load' in metafunc.funcargnames:
        pys = 'py2', 'py3'
        for dump in pys:
            for load in pys:
                param = (dump, load)
                conversion = '%s to %s'%param
                if 'repr' not in metafunc.funcargnames:
                    metafunc.addcall(id=conversion, param=param)
                else:
                    for tp, repr in simple_tests.items():
                        metafunc.addcall(
                            id='%s:%s'%(tp, conversion),
                            param=param,
                            funcargs={'tp_name':tp, 'repr':repr},
                            )

simple_tests = {
#   type: expected before/after repr
    'int': '4',
    'float':'3.25',
    'list': '[1, 2, 3]',
    'tuple': '(1, 2, 3)',
    'dict': '{(1, 2, 3): 32}',
}

def test_simple(tp_name, repr, dump, load):
    p = dump(repr)
    tp , v = load(p)
    assert tp == tp_name
    assert v == repr

def test_set(py2, py3):
    for dump in py2.dump, py3.dump:
        p = dump("set((1, 2, 3))")
        tp, v = py2.load(p)
        assert tp == "set"
        assert v == "set([1, 2, 3])"
        tp, v = py3.load(p)
        assert tp == "set"
        assert v == "{1, 2, 3}"
        p = dump("set()")
        tp, v = py2.load(p)
        assert tp == "set"
        assert v == "set([])"
        tp, v = py3.load(p)
        assert tp == "set"
        assert v == "set()"

def test_frozenset(py2, py3):
    for dump in py2.dump, py3.dump:
        p = dump("frozenset((1, 2, 3))")
        tp, v = py2.load(p)
        assert tp == "frozenset"
        assert v == "frozenset([1, 2, 3])"
        tp, v = py3.load(p)
        assert tp == "frozenset"
        assert v == "frozenset({1, 2, 3})"
        p = dump("frozenset()")
        tp, v = py2.load(p)
        assert tp == "frozenset"
        assert v == "frozenset([])"
        tp, v = py3.load(p)
        assert tp == "frozenset"
        assert v == "frozenset()"

def test_long(py2, py3):
    really_big = "9223372036854775807324234"
    p = py2.dump(really_big)
    tp, v = py2.load(p)
    assert tp == "long"
    assert v == really_big + "L"
    tp, v = py3.load(p)
    assert tp == "int"
    assert v == really_big
    p = py3.dump(really_big)
    tp, v == py3.load(p)
    assert tp == "int"
    assert v == really_big
    tp, v = py2.load(p)
    assert tp == "long"
    assert v == really_big + "L"

def test_small_long(py2, py3):
    p = py2.dump("123L")
    tp, s = py2.load(p)
    assert s == "123L"
    tp, s = py3.load(p)
    assert s == "123"

@py.test.mark.xfail
# I'm not sure if we need the complexity.
def test_recursive_list(py2, py3):
    l = [1, 2, 3]
    l.append(l)
    p = py2.dump(l)
    tp, rep = py2.load(l)
    assert tp == "list"

def test_bytes(py2, py3):
    p = py3.dump("b'hi'")
    tp, v = py2.load(p)
    assert tp == "str"
    assert v == "'hi'"
    tp, v = py3.load(p)
    assert tp == "bytes"
    assert v == "b'hi'"

def test_str(py2, py3):
    p = py2.dump("'xyz'")
    tp, s = py2.load(p)
    assert tp == "str"
    assert s == "'xyz'"
    tp, s = py3.load(p, "py2str_as_py3str=True")
    assert tp == "str"
    assert s == "'xyz'"
    tp, s = py3.load(p, "py2str_as_py3str=False")
    assert s == "b'xyz'"
    assert tp == "bytes"

def test_unicode(py2, py3):
    p = py2.dump("u'hi'")
    tp, s = py2.load(p)
    assert tp == "unicode"
    assert s == "u'hi'"
    tp, s = py3.load(p)
    assert tp == "str"
    assert s == "'hi'"
    p = py3.dump("'hi'")
    tp, s = py3.load(p)
    assert tp == "str"
    assert s == "'hi'"
    tp, s = py2.load(p)
    assert tp == "unicode" # depends on unserialization defaults
    assert s == "u'hi'"

def test_bool(py2, py3):
    p = py2.dump("True")
    tp, s = py2.load(p)
    assert tp == "bool"
    assert s == "True"
    tp, s = py3.load(p)
    assert s == "True"
    assert tp == "bool"
    p = py2.dump("False")
    tp, s = py2.load(p)
    assert s == "False"

def test_none(py2, py3):
    p = py2.dump("None")
    tp, s = py2.load(p)
    assert s == "None"
    tp, s = py3.load(p)
    assert s == "None"

def test_tuple_nested_with_empty_in_between(py2):
    p = py2.dump("(1, (), 3)")
    tp, s = py2.load(p)
    assert tp == 'tuple'
    assert s == "(1, (), 3)"
<commit_msg>use popen.communicate()<commit_after># -*- coding: utf-8 -*-
import sys
import os
import tempfile
import subprocess
import py
import execnet
from execnet import gateway_base as serializer

def _find_version(suffix=""):
    name = "python" + suffix
    executable = py.path.local.sysfind(name)
    if executable is None:
        if sys.platform == "win32" and suffix == "3":
            for name in ('python31', 'python30'):
                executable = py.path.local(r"c:\\%s\python.exe" % (name,))
                if executable.check():
                    return executable
        py.test.skip("can't find a %r executable" % (name,))
    return executable

def setup_module(mod):
    mod.TEMPDIR = py.path.local(tempfile.mkdtemp())
    if sys.version_info > (3, 0):
        mod._py3_wrapper = PythonWrapper(py.path.local(sys.executable))
        mod._py2_wrapper = PythonWrapper(_find_version())
    else:
        mod._py3_wrapper = PythonWrapper(_find_version("3"))
        mod._py2_wrapper = PythonWrapper(py.path.local(sys.executable))

def teardown_module(mod):
    TEMPDIR.remove(True)

pyimportdir = str(py.path.local(execnet.__file__).dirpath().dirpath())
class PythonWrapper(object):

    def __init__(self, executable):
        self.executable = executable

    def dump(self, obj_rep):
        script_file = TEMPDIR.join("dump.py")
        script_file.write("""
import sys
sys.path.insert(0, %r)
from execnet import gateway_base as serializer
if sys.version_info > (3, 0): # Need binary output
    sys.stdout = sys.stdout.detach()
saver = serializer.Serializer(sys.stdout)
saver.save(%s)""" % (pyimportdir, obj_rep,))
        popen = subprocess.Popen([str(self.executable), str(script_file)],
                                 stdin=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 stdout=subprocess.PIPE)
        stdout, stderr = popen.communicate()
        ret = popen.returncode
        if ret:
            raise py.process.cmdexec.Error(ret, ret, str(self.executable),
                                           stdout, stderr)
        return stdout

    def load(self, data, option_args="__class__"):
        script_file = TEMPDIR.join("load.py")
        script_file.write(r"""
import sys
sys.path.insert(0, %r)
from execnet import gateway_base as serializer
if sys.version_info > (3, 0):
    sys.stdin = sys.stdin.detach()
loader = serializer.Unserializer(sys.stdin)
loader.%s
obj = loader.load()
sys.stdout.write(type(obj).__name__ + "\n")
sys.stdout.write(repr(obj))""" % (pyimportdir, option_args,))
        popen = subprocess.Popen([str(self.executable), str(script_file)],
                                 stdin=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 stdout=subprocess.PIPE)
        stdout, stderr = popen.communicate(data)
        ret = popen.returncode
        if ret:
            raise py.process.cmdexec.Error(ret, ret, str(self.executable),
                                           stdout, stderr)
        return [s.decode("ascii") for s in stdout.splitlines()]

    def __repr__(self):
        return "<PythonWrapper for %s>" % (self.executable,)

def pytest_funcarg__py2(request):
    return _py2_wrapper

def pytest_funcarg__py3(request):
    return _py3_wrapper

def pytest_funcarg__dump(request):
    py_dump = request.getfuncargvalue(request.param[0])
    return py_dump.dump

def pytest_funcarg__load(request):
    py_dump = request.getfuncargvalue(request.param[1])
    return py_dump.load

def pytest_generate_tests(metafunc):
    if 'dump' in metafunc.funcargnames and 'load' in metafunc.funcargnames:
        pys = 'py2', 'py3'
        for dump in pys:
            for load in pys:
                param = (dump, load)
                conversion = '%s to %s'%param
                if 'repr' not in metafunc.funcargnames:
                    metafunc.addcall(id=conversion, param=param)
                else:
                    for tp, repr in simple_tests.items():
                        metafunc.addcall(
                            id='%s:%s'%(tp, conversion),
                            param=param,
                            funcargs={'tp_name':tp, 'repr':repr},
                            )

simple_tests = {
#   type: expected before/after repr
    'int': '4',
    'float':'3.25',
    'list': '[1, 2, 3]',
    'tuple': '(1, 2, 3)',
    'dict': '{(1, 2, 3): 32}',
}

def test_simple(tp_name, repr, dump, load):
    p = dump(repr)
    tp , v = load(p)
    assert tp == tp_name
    assert v == repr

def test_set(py2, py3):
    for dump in py2.dump, py3.dump:
        p = dump("set((1, 2, 3))")
        tp, v = py2.load(p)
        assert tp == "set"
        assert v == "set([1, 2, 3])"
        tp, v = py3.load(p)
        assert tp == "set"
        assert v == "{1, 2, 3}"
        p = dump("set()")
        tp, v = py2.load(p)
        assert tp == "set"
        assert v == "set([])"
        tp, v = py3.load(p)
        assert tp == "set"
        assert v == "set()"

def test_frozenset(py2, py3):
    for dump in py2.dump, py3.dump:
        p = dump("frozenset((1, 2, 3))")
        tp, v = py2.load(p)
        assert tp == "frozenset"
        assert v == "frozenset([1, 2, 3])"
        tp, v = py3.load(p)
        assert tp == "frozenset"
        assert v == "frozenset({1, 2, 3})"
        p = dump("frozenset()")
        tp, v = py2.load(p)
        assert tp == "frozenset"
        assert v == "frozenset([])"
        tp, v = py3.load(p)
        assert tp == "frozenset"
        assert v == "frozenset()"

def test_long(py2, py3):
    really_big = "9223372036854775807324234"
    p = py2.dump(really_big)
    tp, v = py2.load(p)
    assert tp == "long"
    assert v == really_big + "L"
    tp, v = py3.load(p)
    assert tp == "int"
    assert v == really_big
    p = py3.dump(really_big)
    tp, v == py3.load(p)
    assert tp == "int"
    assert v == really_big
    tp, v = py2.load(p)
    assert tp == "long"
    assert v == really_big + "L"

def test_small_long(py2, py3):
    p = py2.dump("123L")
    tp, s = py2.load(p)
    assert s == "123L"
    tp, s = py3.load(p)
    assert s == "123"

@py.test.mark.xfail
# I'm not sure if we need the complexity.
def test_recursive_list(py2, py3):
    l = [1, 2, 3]
    l.append(l)
    p = py2.dump(l)
    tp, rep = py2.load(l)
    assert tp == "list"

def test_bytes(py2, py3):
    p = py3.dump("b'hi'")
    tp, v = py2.load(p)
    assert tp == "str"
    assert v == "'hi'"
    tp, v = py3.load(p)
    assert tp == "bytes"
    assert v == "b'hi'"

def test_str(py2, py3):
    p = py2.dump("'xyz'")
    tp, s = py2.load(p)
    assert tp == "str"
    assert s == "'xyz'"
    tp, s = py3.load(p, "py2str_as_py3str=True")
    assert tp == "str"
    assert s == "'xyz'"
    tp, s = py3.load(p, "py2str_as_py3str=False")
    assert s == "b'xyz'"
    assert tp == "bytes"

def test_unicode(py2, py3):
    p = py2.dump("u'hi'")
    tp, s = py2.load(p)
    assert tp == "unicode"
    assert s == "u'hi'"
    tp, s = py3.load(p)
    assert tp == "str"
    assert s == "'hi'"
    p = py3.dump("'hi'")
    tp, s = py3.load(p)
    assert tp == "str"
    assert s == "'hi'"
    tp, s = py2.load(p)
    assert tp == "unicode" # depends on unserialization defaults
    assert s == "u'hi'"

def test_bool(py2, py3):
    p = py2.dump("True")
    tp, s = py2.load(p)
    assert tp == "bool"
    assert s == "True"
    tp, s = py3.load(p)
    assert s == "True"
    assert tp == "bool"
    p = py2.dump("False")
    tp, s = py2.load(p)
    assert s == "False"

def test_none(py2, py3):
    p = py2.dump("None")
    tp, s = py2.load(p)
    assert s == "None"
    tp, s = py3.load(p)
    assert s == "None"

def test_tuple_nested_with_empty_in_between(py2):
    p = py2.dump("(1, (), 3)")
    tp, s = py2.load(p)
    assert tp == 'tuple'
    assert s == "(1, (), 3)"
<|endoftext|>


<commit_before>from os import listdir, rename, walk
from os.path import isfile, join, basename, abspath, dirname
import sys
import hashlib

def sha1OfFile(filepath):
	filepath = abspath(filepath)
	with open(filepath, 'rb') as f:
		return hashlib.sha1(f.read()).hexdigest()

if sys.argv[1][len(sys.argv[1]) - 1] == '/':
	sys.argv[1] = sys.argv[1][:-1]

if sys.argv[2][len(sys.argv[2]) - 1] == '/':
	sys.argv[2] = sys.argv[2][:-1]

original_files = {}

for root, subFolders, files in walk(sys.argv[1]):
	for f in files:
		p = join(root, f)
		f = abspath(p)
		original_files[sha1OfFile(f)] = basename(f)

found = 0;
not_found = 0;

for root, subFolders, files in walk(sys.argv[2]):
	for f in files:
		p = join(root, f)
		f = abspath(p)

		if isfile(join(sys.argv[2], f)):
			if sha1OfFile(f) in original_files:
				found += 1

				rename(f, dirname(abspath(f)) + '/' + original_files[sha1OfFile(f)])
			else:
				not_found = 1

print 'Total original files: ' + str(len(original_files))
print 'Total recovered files found: ' + str(found)
print 'Total not recovered files found: ' + str(not_found)<commit_msg>Fixed not found files counter<commit_after>from os import listdir, rename, walk
from os.path import isfile, join, basename, abspath, dirname

import sys
import hashlib

def sha1OfFile(filepath):
	filepath = abspath(filepath)
	with open(filepath, 'rb') as f:
		return hashlib.sha1(f.read()).hexdigest()

if sys.argv[1][len(sys.argv[1]) - 1] == '/':
	sys.argv[1] = sys.argv[1][:-1]

if sys.argv[2][len(sys.argv[2]) - 1] == '/':
	sys.argv[2] = sys.argv[2][:-1]

original_files = {}

for root, subFolders, files in walk(sys.argv[1]):
	for f in files:
		p = join(root, f)
		f = abspath(p)
		original_files[sha1OfFile(f)] = basename(f)

found = 0;
not_found = 0;

for root, subFolders, files in walk(sys.argv[2]):
	for f in files:
		p = join(root, f)
		f = abspath(p)

		if isfile(join(sys.argv[2], f)):
			if sha1OfFile(f) in original_files:
				found += 1

				rename(f, dirname(abspath(f)) + '/' + original_files[sha1OfFile(f)])
			else:
				not_found += 1

print 'Total original files: ' + str(len(original_files))
print 'Total recovered files found: ' + str(found)
print 'Total not recovered files found: ' + str(not_found)
<|endoftext|>


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

import sys, os, re

# Pass date and base path to main() from airflow
def main(base_path):

  # Default to "."
  try: base_path
  except NameError: base_path = "."
  if not base_path:
    base_path = "."

  APP_NAME = "train_spark_mllib_model.py"

  # If there is no SparkSession, create the environment
  try:
    sc and spark
  except (NameError, UnboundLocalError) as e:
    import findspark
    findspark.init()
    import pyspark
    import pyspark.sql

    sc = pyspark.SparkContext()
    spark = pyspark.sql.SparkSession(sc).builder.appName(APP_NAME).getOrCreate()

  #
  # {
  #   "ArrDelay":5.0,"CRSArrTime":"2015-12-31T03:20:00.000-08:00","CRSDepTime":"2015-12-31T03:05:00.000-08:00",
  #   "Carrier":"WN","DayOfMonth":31,"DayOfWeek":4,"DayOfYear":365,"DepDelay":14.0,"Dest":"SAN","Distance":368.0,
  #   "FlightDate":"2015-12-30T16:00:00.000-08:00","FlightNum":"6109","Origin":"TUS"
  # }
  #
  from pyspark.sql.types import StringType, IntegerType, FloatType, DoubleType, DateType, TimestampType
  from pyspark.sql.types import StructType, StructField
  from pyspark.sql.functions import udf

  schema = StructType([
    StructField("ArrDelay", DoubleType(), True),     # "ArrDelay":5.0
    StructField("CRSArrTime", TimestampType(), True),    # "CRSArrTime":"2015-12-31T03:20:00.000-08:00"
    StructField("CRSDepTime", TimestampType(), True),    # "CRSDepTime":"2015-12-31T03:05:00.000-08:00"
    StructField("Carrier", StringType(), True),     # "Carrier":"WN"
    StructField("DayOfMonth", IntegerType(), True), # "DayOfMonth":31
    StructField("DayOfWeek", IntegerType(), True),  # "DayOfWeek":4
    StructField("DayOfYear", IntegerType(), True),  # "DayOfYear":365
    StructField("DepDelay", DoubleType(), True),     # "DepDelay":14.0
    StructField("Dest", StringType(), True),        # "Dest":"SAN"
    StructField("Distance", DoubleType(), True),     # "Distance":368.0
    StructField("FlightDate", DateType(), True),    # "FlightDate":"2015-12-30T16:00:00.000-08:00"
    StructField("FlightNum", StringType(), True),   # "FlightNum":"6109"
    StructField("Origin", StringType(), True),      # "Origin":"TUS"
  ])

  input_path = "{}/data/simple_flight_delay_features.jsonl.bz2".format(
    base_path
  )
  features = spark.read.json(input_path, schema=schema)
  features.first()

  #
  # Check for nulls in features before using Spark ML
  #
  null_counts = [(column, features.where(features[column].isNull()).count()) for column in features.columns]
  cols_with_nulls = filter(lambda x: x[1] > 0, null_counts)
  print(list(cols_with_nulls))

  #
  # Add a Route variable to replace FlightNum
  #
  from pyspark.sql.functions import lit, concat
  features_with_route = features.withColumn(
    'Route',
    concat(
      features.Origin,
      lit('-'),
      features.Dest
    )
  )
  features_with_route.show(6)

  #
  # Use pysmark.ml.feature.Bucketizer to bucketize ArrDelay into on-time, slightly late, very late (0, 1, 2)
  #
  from pyspark.ml.feature import Bucketizer

  # Setup the Bucketizer
  splits = [-float("inf"), -15.0, 0, 30.0, float("inf")]
  arrival_bucketizer = Bucketizer(
    splits=splits,
    inputCol="ArrDelay",
    outputCol="ArrDelayBucket"
  )

  # Save the bucketizer
  arrival_bucketizer_path = "{}/models/arrival_bucketizer_2.0.bin".format(base_path)
  arrival_bucketizer.write().overwrite().save(arrival_bucketizer_path)

  # Setup the Departure Bucketizer for other examples
  splits = [-float("inf"), -15.0, 0, 30.0, float("inf")]
  departure_bucketizer = Bucketizer(
    splits=splits,
    inputCol="DepDelay",
    outputCol="DepDelayBucket"
  )

  # Save the departure bucketizer
  departure_bucketizer_path = "{}/models/arrival_bucketizer_2.0.bin".format(base_path)
  departure_bucketizer.write().overwrite().save(departure_bucketizer_path)

  # Apply the bucketizer
  ml_bucketized_features = arrival_bucketizer.transform(features_with_route)
  ml_bucketized_features.select("DepDelay", "DepDelayBucket").show()

  #
  # Extract features tools in with pyspark.ml.feature
  #
  from pyspark.ml.feature import StringIndexer, VectorAssembler

  # Turn category fields into indexes
  for column in ["Carrier", "Origin", "Dest", "Route"]:
    string_indexer = StringIndexer(
      inputCol=column,
      outputCol=column + "_index"
    )

    string_indexer_model = string_indexer.fit(ml_bucketized_features)
    ml_bucketized_features = string_indexer_model.transform(ml_bucketized_features)

    # Drop the original column
    ml_bucketized_features = ml_bucketized_features.drop(column)

    # Save the pipeline model
    string_indexer_output_path = "{}/models/string_indexer_model_{}.bin".format(
      base_path,
      column
    )
    string_indexer_model.write().overwrite().save(string_indexer_output_path)

  # Combine continuous, numeric fields with indexes of nominal ones
  # ...into one feature vector
  numeric_columns = [
    "DepDelay", "Distance",
    "DayOfMonth", "DayOfWeek",
    "DayOfYear"]
  index_columns = ["Carrier_index", "Origin_index",
                   "Dest_index", "Route_index"]
  vector_assembler = VectorAssembler(
    inputCols=numeric_columns + index_columns,
    outputCol="Features_vec"
  )
  final_vectorized_features = vector_assembler.transform(ml_bucketized_features)

  # Save the numeric vector assembler
  vector_assembler_path = "{}/models/numeric_vector_assembler.bin".format(base_path)
  vector_assembler.write().overwrite().save(vector_assembler_path)

  # Drop the index columns
  for column in index_columns:
    final_vectorized_features = final_vectorized_features.drop(column)

  # Inspect the finalized features
  final_vectorized_features.show()

  # Instantiate and fit random forest classifier on all the data
  from pyspark.ml.classification import RandomForestClassifier
  rfc = RandomForestClassifier(
    featuresCol="Features_vec",
    labelCol="ArrDelayBucket",
    predictionCol="Prediction",
    maxBins=4657,
    maxMemoryInMB=1024
  )
  model = rfc.fit(final_vectorized_features)

  # Save the new model over the old one
  model_output_path = "{}/models/spark_random_forest_classifier.flight_delays.5.0.bin".format(
    base_path
  )
  model.write().overwrite().save(model_output_path)

  # Evaluate model using test data
  predictions = model.transform(final_vectorized_features)

  from pyspark.ml.evaluation import MulticlassClassificationEvaluator
  evaluator = MulticlassClassificationEvaluator(
    predictionCol="Prediction",
    labelCol="ArrDelayBucket",
    metricName="accuracy"
  )
  accuracy = evaluator.evaluate(predictions)
  print("Accuracy = {}".format(accuracy))

  # Check the distribution of predictions
  predictions.groupBy("Prediction").count().show()

  # Check a sample
  predictions.sample(False, 0.001, 18).orderBy("CRSDepTime").show(6)

if __name__ == "__main__":
  main(sys.argv[1])
<commit_msg>Fixed bug<commit_after># !/usr/bin/env python

import sys, os, re

# Pass date and base path to main() from airflow
def main(base_path):

  # Default to "."
  try: base_path
  except NameError: base_path = "."
  if not base_path:
    base_path = "."

  APP_NAME = "train_spark_mllib_model.py"

  # If there is no SparkSession, create the environment
  try:
    sc and spark
  except (NameError, UnboundLocalError) as e:
    import findspark
    findspark.init()
    import pyspark
    import pyspark.sql

    sc = pyspark.SparkContext()
    spark = pyspark.sql.SparkSession(sc).builder.appName(APP_NAME).getOrCreate()

  #
  # {
  #   "ArrDelay":5.0,"CRSArrTime":"2015-12-31T03:20:00.000-08:00","CRSDepTime":"2015-12-31T03:05:00.000-08:00",
  #   "Carrier":"WN","DayOfMonth":31,"DayOfWeek":4,"DayOfYear":365,"DepDelay":14.0,"Dest":"SAN","Distance":368.0,
  #   "FlightDate":"2015-12-30T16:00:00.000-08:00","FlightNum":"6109","Origin":"TUS"
  # }
  #
  from pyspark.sql.types import StringType, IntegerType, FloatType, DoubleType, DateType, TimestampType
  from pyspark.sql.types import StructType, StructField
  from pyspark.sql.functions import udf

  schema = StructType([
    StructField("ArrDelay", DoubleType(), True),     # "ArrDelay":5.0
    StructField("CRSArrTime", TimestampType(), True),    # "CRSArrTime":"2015-12-31T03:20:00.000-08:00"
    StructField("CRSDepTime", TimestampType(), True),    # "CRSDepTime":"2015-12-31T03:05:00.000-08:00"
    StructField("Carrier", StringType(), True),     # "Carrier":"WN"
    StructField("DayOfMonth", IntegerType(), True), # "DayOfMonth":31
    StructField("DayOfWeek", IntegerType(), True),  # "DayOfWeek":4
    StructField("DayOfYear", IntegerType(), True),  # "DayOfYear":365
    StructField("DepDelay", DoubleType(), True),     # "DepDelay":14.0
    StructField("Dest", StringType(), True),        # "Dest":"SAN"
    StructField("Distance", DoubleType(), True),     # "Distance":368.0
    StructField("FlightDate", DateType(), True),    # "FlightDate":"2015-12-30T16:00:00.000-08:00"
    StructField("FlightNum", StringType(), True),   # "FlightNum":"6109"
    StructField("Origin", StringType(), True),      # "Origin":"TUS"
  ])

  input_path = "{}/data/simple_flight_delay_features.jsonl.bz2".format(
    base_path
  )
  features = spark.read.json(input_path, schema=schema)
  features.first()

  #
  # Check for nulls in features before using Spark ML
  #
  null_counts = [(column, features.where(features[column].isNull()).count()) for column in features.columns]
  cols_with_nulls = filter(lambda x: x[1] > 0, null_counts)
  print(list(cols_with_nulls))

  #
  # Add a Route variable to replace FlightNum
  #
  from pyspark.sql.functions import lit, concat
  features_with_route = features.withColumn(
    'Route',
    concat(
      features.Origin,
      lit('-'),
      features.Dest
    )
  )
  features_with_route.show(6)

  #
  # Use pysmark.ml.feature.Bucketizer to bucketize ArrDelay into on-time, slightly late, very late (0, 1, 2)
  #
  from pyspark.ml.feature import Bucketizer

  # Setup the Bucketizer
  splits = [-float("inf"), -15.0, 0, 30.0, float("inf")]
  arrival_bucketizer = Bucketizer(
    splits=splits,
    inputCol="ArrDelay",
    outputCol="ArrDelayBucket"
  )

  # Save the bucketizer
  arrival_bucketizer_path = "{}/models/arrival_bucketizer_2.0.bin".format(base_path)
  arrival_bucketizer.write().overwrite().save(arrival_bucketizer_path)

  # Setup the Departure Bucketizer for other examples
  splits = [-float("inf"), -15.0, 0, 30.0, float("inf")]
  departure_bucketizer = Bucketizer(
    splits=splits,
    inputCol="DepDelay",
    outputCol="DepDelayBucket"
  )

  # Save the departure bucketizer
  departure_bucketizer_path = "{}/models/arrival_bucketizer_2.0.bin".format(base_path)
  departure_bucketizer.write().overwrite().save(departure_bucketizer_path)

  # Apply the bucketizer
  ml_bucketized_features = arrival_bucketizer.transform(features_with_route)
  ml_bucketized_features.select("ArrDelay", "ArrDelayBucket").show()

  #
  # Extract features tools in with pyspark.ml.feature
  #
  from pyspark.ml.feature import StringIndexer, VectorAssembler

  # Turn category fields into indexes
  for column in ["Carrier", "Origin", "Dest", "Route"]:
    string_indexer = StringIndexer(
      inputCol=column,
      outputCol=column + "_index"
    )

    string_indexer_model = string_indexer.fit(ml_bucketized_features)
    ml_bucketized_features = string_indexer_model.transform(ml_bucketized_features)

    # Drop the original column
    ml_bucketized_features = ml_bucketized_features.drop(column)

    # Save the pipeline model
    string_indexer_output_path = "{}/models/string_indexer_model_{}.bin".format(
      base_path,
      column
    )
    string_indexer_model.write().overwrite().save(string_indexer_output_path)

  # Combine continuous, numeric fields with indexes of nominal ones
  # ...into one feature vector
  numeric_columns = [
    "DepDelay", "Distance",
    "DayOfMonth", "DayOfWeek",
    "DayOfYear"]
  index_columns = ["Carrier_index", "Origin_index",
                   "Dest_index", "Route_index"]
  vector_assembler = VectorAssembler(
    inputCols=numeric_columns + index_columns,
    outputCol="Features_vec"
  )
  final_vectorized_features = vector_assembler.transform(ml_bucketized_features)

  # Save the numeric vector assembler
  vector_assembler_path = "{}/models/numeric_vector_assembler.bin".format(base_path)
  vector_assembler.write().overwrite().save(vector_assembler_path)

  # Drop the index columns
  for column in index_columns:
    final_vectorized_features = final_vectorized_features.drop(column)

  # Inspect the finalized features
  final_vectorized_features.show()

  # Instantiate and fit random forest classifier on all the data
  from pyspark.ml.classification import RandomForestClassifier
  rfc = RandomForestClassifier(
    featuresCol="Features_vec",
    labelCol="ArrDelayBucket",
    predictionCol="Prediction",
    maxBins=4657,
    maxMemoryInMB=1024
  )
  model = rfc.fit(final_vectorized_features)

  # Save the new model over the old one
  model_output_path = "{}/models/spark_random_forest_classifier.flight_delays.5.0.bin".format(
    base_path
  )
  model.write().overwrite().save(model_output_path)

  # Evaluate model using test data
  predictions = model.transform(final_vectorized_features)

  from pyspark.ml.evaluation import MulticlassClassificationEvaluator
  evaluator = MulticlassClassificationEvaluator(
    predictionCol="Prediction",
    labelCol="ArrDelayBucket",
    metricName="accuracy"
  )
  accuracy = evaluator.evaluate(predictions)
  print("Accuracy = {}".format(accuracy))

  # Check the distribution of predictions
  predictions.groupBy("Prediction").count().show()

  # Check a sample
  predictions.sample(False, 0.001, 18).orderBy("CRSDepTime").show(6)

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