<commit_before>import os
import time

import astropy.units as u
from astropy.coordinates import EarthLocation

from . import scheduler as scheduler
from .utils.indi import PanIndiServer

from .utils import *

@has_logger
class Observatory(object):

    """
    Main Observatory class
    """

    def __init__(self, config=None, *args, **kwargs):
        """
        Starts up the observatory. Reads config file, sets up location,
        dates, mount, cameras, and weather station
        """
        assert config is not None, self.logger.warning("Config not set for observatory")
        self.config = config

        self.logger.info('Initializing observatory')

        # Setup information about site location
        self.logger.info('\t Setting up location')
        self._setup_location()

        self.indi_server = PanIndiServer()

        self.logger.info('\t Setting up mount')
        self.mount = None
        self._create_mount()

        self.logger.info('\t Setting up cameras')
        self.cameras = list()
        self._create_cameras()

        self.logger.info('\t Setting up scheduler')
        self.scheduler = None
        self._create_scheduler()

        self.logger.info('Observatory')

##################################################################################################
# Methods
##################################################################################################

    def get_target(self):
        """ Gets the next target from the scheduler """

        self.logger.debug("Getting target for observatory")
        target = self.scheduler.get_target()
        self.logger.debug("Got target for observatory: {}".format(target))

        return target

    def track_target(self, target, hours=2.0):
        """ Track a target for set amount of time.

        This is a utility method that will track a given `target` for a certain number of `hours`.

        WARNING:
            This is a blocking method! It is a utility method only.

        Args:
            target(SkyCoord):   An astropy.coordinates.SkyCoord.
            hours(float):       Number of hours to track for.
        """
        self.logger.info("Tracking target {} for {} hours".format(target, hours))

        self.mount.set_target_coordinates(target)
        self.mount.slew_to_target()

        self.logger.info("Slewing to {}".format(target))
        while self.mount.is_slewing:
            time.sleep(5)

        self.logger.info("Tracking target. Sleeping for {} hours".format(hours))
        time.sleep(hours * 60 * 60)
        self.logger.info("I just finished tracking {}".format(target))

##################################################################################################
# Private Methods
##################################################################################################

    def _setup_location(self):
        """
        Sets up the site and location details for the observatory

        Note:
            These items are read from the 'site' config directive and include:
                * name
                * latitude
                * longitude
                * timezone
                * presseure
                * elevation
                * horizon

        """
        self.logger.info('Setting up site details of observatory')

        if 'location' in self.config:
            config_site = self.config.get('location')

            name = config_site.get('name', 'Nameless Location')

            latitude = config_site.get('latitude') * u.degree
            longitude = config_site.get('longitude') * u.degree

            timezone = config_site.get('timezone')
            utc_offset = config_site.get('offset')

            pressure = config_site.get('pressure', 0.680) * u.bar
            elevation = config_site.get('elevation', 0) * u.meter
            horizon = config_site.get('horizon', 0) * u.degree

            self.location = {
                'name': name,
                'latitude': latitude,
                'longitude': longitude,
                'elevation': elevation,
                'timezone': timezone,
                'utc_offset': utc_offset,
                'pressure': pressure,
                'horizon': horizon
            }
            self.logger.debug("location set: {}".format(self.location))
            self.logger.debug("setting earth_location: {}".format(self.location))
            # Create an EarthLocation for the mount
            location = EarthLocation(
                lat=self.location.get('latitude'),
                lon=self.location.get('longitude'),
                height=self.location.get('elevation'),
            )
            self.earth_location = location
        else:
            raise error.Error(msg='Bad site information')

    def _create_mount(self, mount_info=None):
        """Creates a mount object.

        Details for the creation of the mount object are held in the
        configuration file or can be passed to the method.

        This method ensures that the proper mount type is loaded.

        Note:
            This does not actually make a serial connection to the mount. To do so,
            call the 'mount.connect()' explicitly.

        Args:
            mount_info (dict):  Configuration items for the mount.

        Returns:
            panoptes.mount:     Returns a sub-class of the mount type
        """
        if mount_info is None:
            mount_info = self.config.get('mount')

        model = mount_info.get('model')

        self.logger.debug('Creating mount: {}'.format(model))

        module = load_module('panoptes.mount.{}'.format(model))

        driver = 'indi_ieq_telescope'

        mount_info['name'] = self.config.get('name')
        mount_info['driver'] = driver
        mount_info['offset'] = self.config.get('location', {})

        self.logger.debug("Loading {} driver for {}.".format(driver, model))
        self.indi_server.load_driver(self.config.get('name'), driver)

        # Make the mount include site information
        self.mount = module.Mount(mount_info, location=self.earth_location)
        self.logger.debug('Mount created')

    def _create_cameras(self, camera_info=None):
        """Creates a camera object(s)

        Creates a camera for each camera item listed in the config. Ensures the
        appropriate camera module is loaded.

        Note:
            This does not actually make a usb connection to the camera. To do so,
            call the 'camear.connect()' explicitly.

        Args:
            camera_info (dict): Configuration items for the cameras.

        Returns:
            list: A list of created camera objects.
        """
        if camera_info is None:
            camera_info = self.config.get('cameras')

        self.logger.debug("Camera config: \n {}".format(camera_info))

        cameras = list()

        for cam_num, camera_config in enumerate(camera_info):
            # Actually import the model of camera
            camera_model = camera_config.get('model')
            # driver = camera_config.get('driver')

            cam_name = 'Cam{}'.format(cam_num)
            # cam_name = 'GPhoto CCD'
            camera_config['name'] = cam_name

            driver = 'indi_gphoto_ccd'

            self.logger.debug('Creating camera: {}'.format(camera_model))

            self.logger.debug("Loading {} driver for {}.".format(driver, cam_name))
            self.indi_server.load_driver(cam_name, driver)

            try:
                module = load_module('panoptes.camera.{}'.format(camera_model))
                cam = module.Camera(camera_config)

                if cam.is_connected:
                    cameras.append(cam)
                else:
                    self.logger.warning("{} not connected. Skipping for now.".format(cam.name))
            except ImportError:
                raise error.NotFound(msg=camera_model)

        self.cameras = cameras
        self.logger.debug("Cameras created.")

    def _create_scheduler(self):
        """ Sets up the scheduler that will be used by the observatory """

        # Read the targets from the file
        targets_path = os.path.join(
            self.config.get('base_dir'),
            'resources/conf_files/targets/',
            self.config.get('targets_file')
        )

        if os.path.exists(targets_path):
            self.logger.debug('Creating scheduler: {}'.format(targets_path))
            self.scheduler = scheduler.Scheduler(targets_file=targets_path, location=self.earth_location)
            self.logger.debug("Scheduler created")
        else:
            self.logger.warning("Targets file does not exist: {}".format(targets_path))
<commit_msg>Add utc_offset<commit_after>import os
import time

import astropy.units as u
from astropy.coordinates import EarthLocation

from . import scheduler as scheduler
from .utils.indi import PanIndiServer

from .utils import *

@has_logger
class Observatory(object):

    """
    Main Observatory class
    """

    def __init__(self, config=None, *args, **kwargs):
        """
        Starts up the observatory. Reads config file, sets up location,
        dates, mount, cameras, and weather station
        """
        assert config is not None, self.logger.warning("Config not set for observatory")
        self.config = config

        self.logger.info('Initializing observatory')

        # Setup information about site location
        self.logger.info('\t Setting up location')
        self._setup_location()

        self.indi_server = PanIndiServer()

        self.logger.info('\t Setting up mount')
        self.mount = None
        self._create_mount()

        self.logger.info('\t Setting up cameras')
        self.cameras = list()
        self._create_cameras()

        self.logger.info('\t Setting up scheduler')
        self.scheduler = None
        self._create_scheduler()

        self.logger.info('Observatory')

##################################################################################################
# Methods
##################################################################################################

    def get_target(self):
        """ Gets the next target from the scheduler """

        self.logger.debug("Getting target for observatory")
        target = self.scheduler.get_target()
        self.logger.debug("Got target for observatory: {}".format(target))

        return target

    def track_target(self, target, hours=2.0):
        """ Track a target for set amount of time.

        This is a utility method that will track a given `target` for a certain number of `hours`.

        WARNING:
            This is a blocking method! It is a utility method only.

        Args:
            target(SkyCoord):   An astropy.coordinates.SkyCoord.
            hours(float):       Number of hours to track for.
        """
        self.logger.info("Tracking target {} for {} hours".format(target, hours))

        self.mount.set_target_coordinates(target)
        self.mount.slew_to_target()

        self.logger.info("Slewing to {}".format(target))
        while self.mount.is_slewing:
            time.sleep(5)

        self.logger.info("Tracking target. Sleeping for {} hours".format(hours))
        time.sleep(hours * 60 * 60)
        self.logger.info("I just finished tracking {}".format(target))

##################################################################################################
# Private Methods
##################################################################################################

    def _setup_location(self):
        """
        Sets up the site and location details for the observatory

        Note:
            These items are read from the 'site' config directive and include:
                * name
                * latitude
                * longitude
                * timezone
                * presseure
                * elevation
                * horizon

        """
        self.logger.info('Setting up site details of observatory')

        if 'location' in self.config:
            config_site = self.config.get('location')

            name = config_site.get('name', 'Nameless Location')

            latitude = config_site.get('latitude') * u.degree
            longitude = config_site.get('longitude') * u.degree

            timezone = config_site.get('timezone')
            utc_offset = config_site.get('utc_offset')

            pressure = config_site.get('pressure', 0.680) * u.bar
            elevation = config_site.get('elevation', 0) * u.meter
            horizon = config_site.get('horizon', 0) * u.degree

            self.location = {
                'name': name,
                'latitude': latitude,
                'longitude': longitude,
                'elevation': elevation,
                'timezone': timezone,
                'utc_offset': utc_offset,
                'pressure': pressure,
                'horizon': horizon
            }
            self.logger.debug("location set: {}".format(self.location))
            self.logger.debug("setting earth_location: {}".format(self.location))
            # Create an EarthLocation for the mount
            location = EarthLocation(
                lat=self.location.get('latitude'),
                lon=self.location.get('longitude'),
                height=self.location.get('elevation'),
            )
            self.earth_location = location
        else:
            raise error.Error(msg='Bad site information')

    def _create_mount(self, mount_info=None):
        """Creates a mount object.

        Details for the creation of the mount object are held in the
        configuration file or can be passed to the method.

        This method ensures that the proper mount type is loaded.

        Note:
            This does not actually make a serial connection to the mount. To do so,
            call the 'mount.connect()' explicitly.

        Args:
            mount_info (dict):  Configuration items for the mount.

        Returns:
            panoptes.mount:     Returns a sub-class of the mount type
        """
        if mount_info is None:
            mount_info = self.config.get('mount')

        model = mount_info.get('model')

        self.logger.debug('Creating mount: {}'.format(model))

        module = load_module('panoptes.mount.{}'.format(model))

        driver = 'indi_ieq_telescope'

        mount_info['name'] = self.config.get('name')
        mount_info['driver'] = driver
        mount_info['utc_offset'] = self.location.get('utc_offset', '0.0')

        self.logger.debug("Loading {} driver for {}.".format(driver, model))
        self.indi_server.load_driver(self.config.get('name'), driver)

        # Make the mount include site information
        self.mount = module.Mount(mount_info, location=self.earth_location)
        self.logger.debug('Mount created')

    def _create_cameras(self, camera_info=None):
        """Creates a camera object(s)

        Creates a camera for each camera item listed in the config. Ensures the
        appropriate camera module is loaded.

        Note:
            This does not actually make a usb connection to the camera. To do so,
            call the 'camear.connect()' explicitly.

        Args:
            camera_info (dict): Configuration items for the cameras.

        Returns:
            list: A list of created camera objects.
        """
        if camera_info is None:
            camera_info = self.config.get('cameras')

        self.logger.debug("Camera config: \n {}".format(camera_info))

        cameras = list()

        for cam_num, camera_config in enumerate(camera_info):
            # Actually import the model of camera
            camera_model = camera_config.get('model')
            # driver = camera_config.get('driver')

            cam_name = 'Cam{}'.format(cam_num)
            # cam_name = 'GPhoto CCD'
            camera_config['name'] = cam_name

            driver = 'indi_gphoto_ccd'

            self.logger.debug('Creating camera: {}'.format(camera_model))

            self.logger.debug("Loading {} driver for {}.".format(driver, cam_name))
            self.indi_server.load_driver(cam_name, driver)

            try:
                module = load_module('panoptes.camera.{}'.format(camera_model))
                cam = module.Camera(camera_config)

                if cam.is_connected:
                    cameras.append(cam)
                else:
                    self.logger.warning("{} not connected. Skipping for now.".format(cam.name))
            except ImportError:
                raise error.NotFound(msg=camera_model)

        self.cameras = cameras
        self.logger.debug("Cameras created.")

    def _create_scheduler(self):
        """ Sets up the scheduler that will be used by the observatory """

        # Read the targets from the file
        targets_path = os.path.join(
            self.config.get('base_dir'),
            'resources/conf_files/targets/',
            self.config.get('targets_file')
        )

        if os.path.exists(targets_path):
            self.logger.debug('Creating scheduler: {}'.format(targets_path))
            self.scheduler = scheduler.Scheduler(targets_file=targets_path, location=self.earth_location)
            self.logger.debug("Scheduler created")
        else:
            self.logger.warning("Targets file does not exist: {}".format(targets_path))
<|endoftext|>


<commit_before>###############################################################################
#
# Copyright 2011 Pants (see AUTHORS.txt)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
###############################################################################

###############################################################################
# Imports
###############################################################################

from pants.network import Connection, Server

###############################################################################
# Constants
###############################################################################

# Telnet commands
IAC  = chr(255) # Interpret As Command
DONT = chr(254) # Don't Perform
DO   = chr(253) # Do Perform
WONT = chr(252) # Won't Perform
WILL = chr(251) # Will Perform
SB   = chr(250) # Subnegotiation Begin
SE   = chr(240) # Subnegotiation End

###############################################################################
# TelnetConnection Class
###############################################################################

class TelnetConnection(Connection):
    """
    A basic implementation of a Telnet connection.

    A TelnetConnection object is capable of identifying and extracting
    Telnet command sequences from incoming data. Upon identifying a
    Telnet command, option or subnegotiation, the connection will call a
    relevant placeholder method. This class should be subclassed to
    provide functionality for individual commands and options.
    """
    def __init__(self, server, socket):
        Connection.__init__(self, server, socket)

        # Initialize Stuff
        self._telnet_data = ""

    ##### Public Event Handlers ###############################################

    def on_command(self, command):
        """
        Placeholder. Called when the connection receives a telnet command,
        such as AYT (Are You There).

        =========  ============
        Argument   Description
        =========  ============
        command    The byte representing the telnet command.
        =========  ============
        """
        pass

    def on_option(self, command, option):
        """
        Placeholder. Called when the connection receives a telnet option
        negotiation sequence, such as IAC WILL ECHO.

        =========  ============
        Argument   Description
        =========  ============
        command    The byte representing the telnet command.
        option     The byte representing the telnet option being negotiated.
        =========  ============
        """
        pass

    def on_subnegotiation(self, option, data):
        """
        Placeholder. Called when the connection receives a subnegotiation
        sequence.

        =========  ============
        Argument   Description
        =========  ============
        option     The byte representing the telnet option for which subnegotiation data has been received.
        data       The received data.
        =========  ============
        """
        pass

    ##### Internal Processing Methods #########################################

    ##### Internal Telnet State Processing ####################################

    def _on_telnet_data(self, data):
        self._telnet_data += data

        while self._telnet_data:
            delimiter = self.read_delimiter

            if delimiter is None:
                data = self._telnet_data
                self._telnet_data = ''
                self._safely_call(self.on_read, data)

            elif isinstance(delimiter, (int,long)):
                if len(self._telnet_data) < delimiter:
                    break
                data = self._telnet_data[:delimiter]
                self._telnet_data = self._telnet_data[delimiter:]
                self._safely_call(self.on_read, data)

            elif isinstance(delimiter, basestring):
                mark = self._telnet_data.find(delimiter)
                if mark == -1:
                    break
                data = self._telnet_data[:mark]
                self._telnet_data = self._telnet_data[mark+len(delimiter):]
                self._safely_call(self.on_read, data)

            else:
                log.warning("Invalid read_delmiter on %s #%d." %
                        (self.__class__.__name__, self.fileno))
                break

            if self._socket is None or not self.connected:
                break

    def _on_telnet_iac(self, data):
        if len(data) < 2:
            return False

        elif data[1] == IAC:
            # It's an escaped IAC byte. Send it to the data buffer.
            self._on_telnet_data(IAC)
            return data[2:]

        elif data[1] in '\xFB\xFC\xFD\xFE':
            if len(data) < 3:
                return False

            self._safely_call(self.on_option, data[1], data[2])
            return data[3:]

        elif data[1] == SB:
            seq = ''
            code = data[2:]
            data = data[3:]
            if not data:
                return False

            while data:
                loc = data.find(IAC)
                if loc == -1:
                    return False

                seq += data[:loc]

                if data[loc+1] == SE:
                    # Match
                    data = data[loc+2:]
                    break

                elif data[loc+1] == IAC:
                    # Escaped
                    seq += IAC
                    data = data[loc+2:]
                    continue

                # Unknown. Skip it.
                data = data[loc+1:]
                if not data:
                    return False

            self._safely_call(self.on_subnegotiation, code, seq)

        # Still here? It must just be a command then. Send it on.
        self._safely_call(self.on_command, data[1])
        return data[2:]

    def _process_recv_buffer(self):
        """
        Completely replace the standard recv buffer processing with a custom
        function for optimal telnet performance.
        """
        while self._recv_buffer:
            loc = self._recv_buffer.find(IAC)

            if loc == -1:
                self._on_telnet_data(self._recv_buffer)
                self._recv_buffer = ''
                break

            elif loc > 0:
                self._on_telnet_data(self._recv_buffer[:loc])
                self._recv_buffer = self._recv_buffer[loc:]

            out = self._on_telnet_iac(self._recv_buffer)
            if out is False:
                break

            self._recv_buffer = out

###############################################################################
# TelnetServer Class
###############################################################################

class TelnetServer(Server):
    """
    A basic implementation of a Telnet server.
    """
    ConnectionClass = TelnetConnection
<commit_msg>Added logging to pants.contrib.telnet.<commit_after>###############################################################################
#
# Copyright 2011 Pants (see AUTHORS.txt)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
###############################################################################

###############################################################################
# Imports
###############################################################################

from pants.network import Connection, Server

###############################################################################
# Logging
###############################################################################

import logging
log = logging.getLogger("pants")

###############################################################################
# Constants
###############################################################################

# Telnet commands
IAC  = chr(255) # Interpret As Command
DONT = chr(254) # Don't Perform
DO   = chr(253) # Do Perform
WONT = chr(252) # Won't Perform
WILL = chr(251) # Will Perform
SB   = chr(250) # Subnegotiation Begin
SE   = chr(240) # Subnegotiation End

###############################################################################
# TelnetConnection Class
###############################################################################

class TelnetConnection(Connection):
    """
    A basic implementation of a Telnet connection.

    A TelnetConnection object is capable of identifying and extracting
    Telnet command sequences from incoming data. Upon identifying a
    Telnet command, option or subnegotiation, the connection will call a
    relevant placeholder method. This class should be subclassed to
    provide functionality for individual commands and options.
    """
    def __init__(self, server, socket):
        Connection.__init__(self, server, socket)

        # Initialize Stuff
        self._telnet_data = ""

    ##### Public Event Handlers ###############################################

    def on_command(self, command):
        """
        Placeholder. Called when the connection receives a telnet command,
        such as AYT (Are You There).

        =========  ============
        Argument   Description
        =========  ============
        command    The byte representing the telnet command.
        =========  ============
        """
        pass

    def on_option(self, command, option):
        """
        Placeholder. Called when the connection receives a telnet option
        negotiation sequence, such as IAC WILL ECHO.

        =========  ============
        Argument   Description
        =========  ============
        command    The byte representing the telnet command.
        option     The byte representing the telnet option being negotiated.
        =========  ============
        """
        pass

    def on_subnegotiation(self, option, data):
        """
        Placeholder. Called when the connection receives a subnegotiation
        sequence.

        =========  ============
        Argument   Description
        =========  ============
        option     The byte representing the telnet option for which subnegotiation data has been received.
        data       The received data.
        =========  ============
        """
        pass

    ##### Internal Processing Methods #########################################

    ##### Internal Telnet State Processing ####################################

    def _on_telnet_data(self, data):
        self._telnet_data += data

        while self._telnet_data:
            delimiter = self.read_delimiter

            if delimiter is None:
                data = self._telnet_data
                self._telnet_data = ''
                self._safely_call(self.on_read, data)

            elif isinstance(delimiter, (int,long)):
                if len(self._telnet_data) < delimiter:
                    break
                data = self._telnet_data[:delimiter]
                self._telnet_data = self._telnet_data[delimiter:]
                self._safely_call(self.on_read, data)

            elif isinstance(delimiter, basestring):
                mark = self._telnet_data.find(delimiter)
                if mark == -1:
                    break
                data = self._telnet_data[:mark]
                self._telnet_data = self._telnet_data[mark+len(delimiter):]
                self._safely_call(self.on_read, data)

            else:
                log.warning("Invalid read_delmiter on %s #%d." %
                        (self.__class__.__name__, self.fileno))
                break

            if self._socket is None or not self.connected:
                break

    def _on_telnet_iac(self, data):
        if len(data) < 2:
            return False

        elif data[1] == IAC:
            # It's an escaped IAC byte. Send it to the data buffer.
            self._on_telnet_data(IAC)
            return data[2:]

        elif data[1] in '\xFB\xFC\xFD\xFE':
            if len(data) < 3:
                return False

            self._safely_call(self.on_option, data[1], data[2])
            return data[3:]

        elif data[1] == SB:
            seq = ''
            code = data[2:]
            data = data[3:]
            if not data:
                return False

            while data:
                loc = data.find(IAC)
                if loc == -1:
                    return False

                seq += data[:loc]

                if data[loc+1] == SE:
                    # Match
                    data = data[loc+2:]
                    break

                elif data[loc+1] == IAC:
                    # Escaped
                    seq += IAC
                    data = data[loc+2:]
                    continue

                # Unknown. Skip it.
                data = data[loc+1:]
                if not data:
                    return False

            self._safely_call(self.on_subnegotiation, code, seq)

        # Still here? It must just be a command then. Send it on.
        self._safely_call(self.on_command, data[1])
        return data[2:]

    def _process_recv_buffer(self):
        """
        Completely replace the standard recv buffer processing with a custom
        function for optimal telnet performance.
        """
        while self._recv_buffer:
            loc = self._recv_buffer.find(IAC)

            if loc == -1:
                self._on_telnet_data(self._recv_buffer)
                self._recv_buffer = ''
                break

            elif loc > 0:
                self._on_telnet_data(self._recv_buffer[:loc])
                self._recv_buffer = self._recv_buffer[loc:]

            out = self._on_telnet_iac(self._recv_buffer)
            if out is False:
                break

            self._recv_buffer = out

###############################################################################
# TelnetServer Class
###############################################################################

class TelnetServer(Server):
    """
    A basic implementation of a Telnet server.
    """
    ConnectionClass = TelnetConnection
<|endoftext|>


<commit_before>from nbclient import NotebookClient
from nbclient.exceptions import CellExecutionError
from traitlets import Bool, Instance

class PapermillNotebookClient(NotebookClient):
    """
    Module containing a  that executes the code cells
    and updates outputs
    """

    log_output = Bool(False).tag(config=True)
    stdout_file = Instance(object, default_value=None).tag(config=True)
    stderr_file = Instance(object, default_value=None).tag(config=True)

    def __init__(self, nb_man, km=None, raise_on_iopub_timeout=True, **kw):
        """Initializes the execution manager.

        Parameters
        ----------
        nb_man : NotebookExecutionManager
            Notebook execution manager wrapper being executed.
        km : KernerlManager (optional)
            Optional kernel manager. If none is provided, a kernel manager will
            be created.
        """
        super().__init__(nb_man.nb, km=km, raise_on_iopub_timeout=raise_on_iopub_timeout, **kw)
        self.nb_man = nb_man

    def execute(self, **kwargs):
        """
        Wraps the parent class process call slightly
        """
        self.reset_execution_trackers()

        with self.setup_kernel(**kwargs):
            self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
            self.papermill_execute_cells()
            info_msg = self._wait_for_reply(self.kc.kernel_info())
            self.nb.metadata['language_info'] = info_msg['content']['language_info']
            self.set_widgets_metadata()

        return self.nb

    def papermill_execute_cells(self):
        """
        This function replaces cell execution with it's own wrapper.

        We are doing this for the following reasons:

        1. Notebooks will stop executing when they encounter a failure but not
           raise a `CellException`. This allows us to save the notebook with the
           traceback even though a `CellExecutionError` was encountered.

        2. We want to write the notebook as cells are executed. We inject our
           logic for that here.

        3. We want to include timing and execution status information with the
           metadata of each cell.
        """
        # Execute each cell and update the output in real time.
        for index, cell in enumerate(self.nb.cells):
            try:
                self.nb_man.cell_start(cell, index)
                self.execute_cell(cell, index)
            except CellExecutionError as ex:
                self.nb_man.cell_exception(self.nb.cells[index], cell_index=index, exception=ex)
                break
            finally:
                self.nb_man.cell_complete(self.nb.cells[index], cell_index=index)

    def log_output_message(self, output):
        """
        Process a given output. May log it in the configured logger and/or write it into
        the configured stdout/stderr files.

        :param output: nbformat.notebooknode.NotebookNode
        :return:
        """
        if output.output_type == "stream":
            content = "".join(output.text)
            if output.name == "stdout":
                if self.log_output:
                    self.log.info(content)
                if self.stdout_file:
                    self.stdout_file.write(content)
                    self.stdout_file.flush()
            elif output.name == "stderr":
                if self.log_output:
                    # In case users want to redirect stderr differently, pipe to warning
                    self.log.warning(content)
                if self.stderr_file:
                    self.stderr_file.write(content)
                    self.stderr_file.flush()
        elif self.log_output and ("data" in output and "text/plain" in output.data):
            self.log.info("".join(output.data['text/plain']))

    def process_message(self, *arg, **kwargs):
        output = super().process_message(*arg, **kwargs)
        self.nb_man.autosave_cell()
        if output and (self.log_output or self.stderr_file or self.stdout_file):
            self.log_output_message(output)
        return output
<commit_msg>Moved to public method for wait_for_reply<commit_after>from nbclient import NotebookClient
from nbclient.exceptions import CellExecutionError
from traitlets import Bool, Instance

class PapermillNotebookClient(NotebookClient):
    """
    Module containing a  that executes the code cells
    and updates outputs
    """

    log_output = Bool(False).tag(config=True)
    stdout_file = Instance(object, default_value=None).tag(config=True)
    stderr_file = Instance(object, default_value=None).tag(config=True)

    def __init__(self, nb_man, km=None, raise_on_iopub_timeout=True, **kw):
        """Initializes the execution manager.

        Parameters
        ----------
        nb_man : NotebookExecutionManager
            Notebook execution manager wrapper being executed.
        km : KernerlManager (optional)
            Optional kernel manager. If none is provided, a kernel manager will
            be created.
        """
        super().__init__(nb_man.nb, km=km, raise_on_iopub_timeout=raise_on_iopub_timeout, **kw)
        self.nb_man = nb_man

    def execute(self, **kwargs):
        """
        Wraps the parent class process call slightly
        """
        self.reset_execution_trackers()

        with self.setup_kernel(**kwargs):
            self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
            self.papermill_execute_cells()
            info_msg = self.wait_for_reply(self.kc.kernel_info())
            self.nb.metadata['language_info'] = info_msg['content']['language_info']
            self.set_widgets_metadata()

        return self.nb

    def papermill_execute_cells(self):
        """
        This function replaces cell execution with it's own wrapper.

        We are doing this for the following reasons:

        1. Notebooks will stop executing when they encounter a failure but not
           raise a `CellException`. This allows us to save the notebook with the
           traceback even though a `CellExecutionError` was encountered.

        2. We want to write the notebook as cells are executed. We inject our
           logic for that here.

        3. We want to include timing and execution status information with the
           metadata of each cell.
        """
        # Execute each cell and update the output in real time.
        for index, cell in enumerate(self.nb.cells):
            try:
                self.nb_man.cell_start(cell, index)
                self.execute_cell(cell, index)
            except CellExecutionError as ex:
                self.nb_man.cell_exception(self.nb.cells[index], cell_index=index, exception=ex)
                break
            finally:
                self.nb_man.cell_complete(self.nb.cells[index], cell_index=index)

    def log_output_message(self, output):
        """
        Process a given output. May log it in the configured logger and/or write it into
        the configured stdout/stderr files.

        :param output: nbformat.notebooknode.NotebookNode
        :return:
        """
        if output.output_type == "stream":
            content = "".join(output.text)
            if output.name == "stdout":
                if self.log_output:
                    self.log.info(content)
                if self.stdout_file:
                    self.stdout_file.write(content)
                    self.stdout_file.flush()
            elif output.name == "stderr":
                if self.log_output:
                    # In case users want to redirect stderr differently, pipe to warning
                    self.log.warning(content)
                if self.stderr_file:
                    self.stderr_file.write(content)
                    self.stderr_file.flush()
        elif self.log_output and ("data" in output and "text/plain" in output.data):
            self.log.info("".join(output.data['text/plain']))

    def process_message(self, *arg, **kwargs):
        output = super().process_message(*arg, **kwargs)
        self.nb_man.autosave_cell()
        if output and (self.log_output or self.stderr_file or self.stdout_file):
            self.log_output_message(output)
        return output
<|endoftext|>


<commit_before># Copyright (C) 2005 John Arbash-Meinel <john@arbash-meinel.com>
# Modified up by: Todd Whiteman <ToddW@ActiveState.com>
#
# This file is part of paramiko.
#
# Paramiko is free software; you can redistribute it and/or modify it under the
# terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 2.1 of the License, or (at your option)
# any later version.
#
# Paramiko 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 Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Paramiko; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.

"""
Functions for communicating with Pageant, the basic windows ssh agent program.
"""

from __future__ import with_statement

import array
import ctypes.wintypes
import platform
import struct
import thread
import threading

from . import _winapi

_AGENT_COPYDATA_ID = 0x804e50ba
_AGENT_MAX_MSGLEN = 8192
# Note: The WM_COPYDATA value is pulled from win32con, as a workaround
# so we do not have to import this huge library just for this one variable.
win32con_WM_COPYDATA = 74

def get_thread_ident():
    try:  # thread.get_ident() exists from Py2.5 to Py2.7.
        return thread.get_ident()
    except AttributeError:  # threading.current_thread().ident exists from Py2.6 up to Py3.4.
        return threading.current_thread().ident

def _get_pageant_window_object():
    return ctypes.windll.user32.FindWindowA('Pageant', 'Pageant')

def can_talk_to_agent():
    """
    Check to see if there is a "Pageant" agent we can talk to.

    This checks both if we have the required libraries (win32all or ctypes)
    and if there is a Pageant currently running.
    """
    return bool(_get_pageant_window_object())

ULONG_PTR = ctypes.c_uint64 if platform.architecture()[0] == '64bit' else ctypes.c_uint32

class COPYDATASTRUCT(ctypes.Structure):
    """
    ctypes implementation of
    http://msdn.microsoft.com/en-us/library/windows/desktop/ms649010%28v=vs.85%29.aspx
    """
    _fields_ = [
        ('num_data', ULONG_PTR),
        ('data_size', ctypes.wintypes.DWORD),
        ('data_loc', ctypes.c_void_p),
    ]

def _query_pageant(msg):
    """
    Communication with the Pageant process is done through a shared
    memory-mapped file.
    """
    hwnd = _get_pageant_window_object()
    if not hwnd:
        # Raise a failure to connect exception, pageant isn't running anymore!
        return None

    # create a name for the mmap
    map_name = 'PageantRequest%08x' % get_thread_ident()

    pymap = _winapi.MemoryMap(map_name, _AGENT_MAX_MSGLEN,
        _winapi.get_security_attributes_for_user(),
        )
    with pymap:
        pymap.write(msg)
        # Create an array buffer containing the mapped filename
        char_buffer = array.array("c", map_name + '\0')
        char_buffer_address, char_buffer_size = char_buffer.buffer_info()
        # Create a string to use for the SendMessage function call
        cds = COPYDATASTRUCT(_AGENT_COPYDATA_ID, char_buffer_size,
            char_buffer_address)

        response = ctypes.windll.user32.SendMessageA(hwnd,
            win32con_WM_COPYDATA, ctypes.sizeof(cds), ctypes.byref(cds))

        if response > 0:
            pymap.seek(0)
            datalen = pymap.read(4)
            retlen = struct.unpack('>I', datalen)[0]
            return datalen + pymap.read(retlen)
        return None

class PageantConnection(object):
    """
    Mock "connection" to an agent which roughly approximates the behavior of
    a unix local-domain socket (as used by Agent).  Requests are sent to the
    pageant daemon via special Windows magick, and responses are buffered back
    for subsequent reads.
    """

    def __init__(self):
        self._response = None

    def send(self, data):
        self._response = _query_pageant(data)

    def recv(self, n):
        if self._response is None:
            return ''
        ret = self._response[:n]
        self._response = self._response[n:]
        if self._response == '':
            self._response = None
        return ret

    def close(self):
        pass
<commit_msg>Try Py2.5 compatibility as last fallback for thread identity.<commit_after># Copyright (C) 2005 John Arbash-Meinel <john@arbash-meinel.com>
# Modified up by: Todd Whiteman <ToddW@ActiveState.com>
#
# This file is part of paramiko.
#
# Paramiko is free software; you can redistribute it and/or modify it under the
# terms of the GNU Lesser General Public License as published by the Free
# Software Foundation; either version 2.1 of the License, or (at your option)
# any later version.
#
# Paramiko 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 Lesser General Public License for more
# details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Paramiko; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.

"""
Functions for communicating with Pageant, the basic windows ssh agent program.
"""

from __future__ import with_statement

import array
import ctypes.wintypes
import platform
import struct
import thread
import threading

from . import _winapi

_AGENT_COPYDATA_ID = 0x804e50ba
_AGENT_MAX_MSGLEN = 8192
# Note: The WM_COPYDATA value is pulled from win32con, as a workaround
# so we do not have to import this huge library just for this one variable.
win32con_WM_COPYDATA = 74

def get_thread_ident():
    # thread.get_ident() exists from Py2.5 to Py2.7.
    # threading.current_thread().ident exists from Py2.6 up to Py3.4.
    try:
        return threading.current_thread().ident
    except AttributeError:
        return thread.get_ident()

def _get_pageant_window_object():
    return ctypes.windll.user32.FindWindowA('Pageant', 'Pageant')

def can_talk_to_agent():
    """
    Check to see if there is a "Pageant" agent we can talk to.

    This checks both if we have the required libraries (win32all or ctypes)
    and if there is a Pageant currently running.
    """
    return bool(_get_pageant_window_object())

ULONG_PTR = ctypes.c_uint64 if platform.architecture()[0] == '64bit' else ctypes.c_uint32

class COPYDATASTRUCT(ctypes.Structure):
    """
    ctypes implementation of
    http://msdn.microsoft.com/en-us/library/windows/desktop/ms649010%28v=vs.85%29.aspx
    """
    _fields_ = [
        ('num_data', ULONG_PTR),
        ('data_size', ctypes.wintypes.DWORD),
        ('data_loc', ctypes.c_void_p),
    ]

def _query_pageant(msg):
    """
    Communication with the Pageant process is done through a shared
    memory-mapped file.
    """
    hwnd = _get_pageant_window_object()
    if not hwnd:
        # Raise a failure to connect exception, pageant isn't running anymore!
        return None

    # create a name for the mmap
    map_name = 'PageantRequest%08x' % get_thread_ident()

    pymap = _winapi.MemoryMap(map_name, _AGENT_MAX_MSGLEN,
        _winapi.get_security_attributes_for_user(),
        )
    with pymap:
        pymap.write(msg)
        # Create an array buffer containing the mapped filename
        char_buffer = array.array("c", map_name + '\0')
        char_buffer_address, char_buffer_size = char_buffer.buffer_info()
        # Create a string to use for the SendMessage function call
        cds = COPYDATASTRUCT(_AGENT_COPYDATA_ID, char_buffer_size,
            char_buffer_address)

        response = ctypes.windll.user32.SendMessageA(hwnd,
            win32con_WM_COPYDATA, ctypes.sizeof(cds), ctypes.byref(cds))

        if response > 0:
            pymap.seek(0)
            datalen = pymap.read(4)
            retlen = struct.unpack('>I', datalen)[0]
            return datalen + pymap.read(retlen)
        return None

class PageantConnection(object):
    """
    Mock "connection" to an agent which roughly approximates the behavior of
    a unix local-domain socket (as used by Agent).  Requests are sent to the
    pageant daemon via special Windows magick, and responses are buffered back
    for subsequent reads.
    """

    def __init__(self):
        self._response = None

    def send(self, data):
        self._response = _query_pageant(data)

    def recv(self, n):
        if self._response is None:
            return ''
        ret = self._response[:n]
        self._response = self._response[n:]
        if self._response == '':
            self._response = None
        return ret

    def close(self):
        pass
<|endoftext|>


<commit_before>import json
from pprint import pprint
import re

json_data=open('crashlacma.20140524-102001.json')

data = json.load(json_data)

# TODO: handle test cases
# testcases:
# hollywood & vine, hollywood and vine
# order of operations: hashtag, img, address, other text.
# hashtag allcaps or lowercase
# uploaded image, link to hosted image

# if image is uploaded via twitter
img_url = data["entities"]["media"][0]["media_url"]
hashtag = data["entities"]["hashtags"][0]["text"]    # won't need this parsed
tweet = data["text"]
# this regex scrapes out URLs. that'll need to change to handle imgs as hosted links
tweet_text = ' '.join(re.sub("(@[A-Za-z0-9]+)|(#[A-Za-z0-9]+)|(\w+:\/\/\S+)"," ",tweet).split())

print(hashtag)
print(img_url)
print(tweet)
print(tweet_text)

json_data.close()<commit_msg>added handling for images sent as url and images uploaded through twitter<commit_after>import json
from pprint import pprint
import re

#json_data=open('sample_http_img.json')
json_data=open('sample_upload_img.json')

data = json.load(json_data)

# TODO: handle test cases
# testcases:
# hollywood & vine, hollywood and vine
# order of operations: hashtag, img, address, other text.
# hashtag allcaps or lowercase
# uploaded image, link to hosted image
# multiple urls? currently hard-coded to only accept the first url seen. probably best this way.

#pprint(data)

tweet = data["text"]
print tweet

if data["entities"].get('media'): # img uploaded via twitter
	###print "img uploaded"
	img_url = data["entities"]["media"][0]["media_url"]
else: # img passed as url
	###print "img as url"
	img_url = re.findall('http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', tweet)[0]

hashtag = data["entities"]["hashtags"][0]["text"] # won't need this parsed

# this regex scrapes out URLs. that'll need to change to handle imgs as hosted links
tweet_text = ' '.join(re.sub("(@[A-Za-z0-9]+)|(#[A-Za-z0-9]+)|(\w+:\/\/\S+)"," ",tweet).split())

print(hashtag)
print(img_url)
print(tweet_text)

json_data.close()<|endoftext|>


<commit_before>#!/usr/bin/python

"""This is the master parsing file. It will take a folder of raw texts, metadata, and marcXML
	and construct the necessary input.txt and jsoncatalog.txt for Presidio to compile."""

#build core files for Presidio database

import re
import string
import os, sys
import subprocess
import xml.etree.ElementTree as ET
#may not need some of these, preserve memory
sys.setrecursionlimit(60000)

# open files, folders, read them
class fileHandlers:

	def __init__(self, directory, directoryAlt):
		self.directory = directory

	#generator that returns next record in specified folder
	def openDir(self):
		for file in os.listdir(self.directory):
	    		if (os.path.isfile(os.path.join(self.directory, file))):
	        	    yield file

	#reads data within a file
	def readItem(self, item):
		file = open((os.path.join(self.directory, item)), 'r')
		data = file.read()
		file.close()
		return data

	#MARC data will be used as a base. If it can't find relevant data in MARC,
	#open file in XML directory instead.
	def readAlternate(self, item, extension):
		#chop off the filetype, replace it
		#note: only possible due to archive.org's impressive uniformity
		item = item[:-9]
		item = item + extension
		file = open((os.path.join(self.directoryAlt, item)), 'r')
		data = file.read()
		file.close()
		return data

#functions that construct the JSON array
class Components:

	def __init__(self):
		self.marcDir = fileHandlers('/data/MARC',
			'/home/will/MHL_download/mhl_meta_xml_files')
		self.xmlDir = fileHandlers('/home/will/MHL_download/mhl_meta_xml_files',
			'/data/MARC')
		self.marcNext = self.marcDir.openDir().next()
		self.xmlNext = self.xmlDir.openDir().next()

		self.marcRoot = ET.fromstring(self.marcDir.readItem(self.marcNext))
		self.xmlRoot = ET.fromstring(self.xmlDir.readItem(self.xmlNext))

#problem: how to align, synchornize two different filetypes
	def getDate(self):
		pass

	def getLibrary(self):
		pass

	def getLanguage(self):
		pass

	def getTitle(self):
		pass

	def getSearchstring(self):
		pass

	def getSubject(self):
		pass

	def getLocation(self):
		pass

def main():
	buildJSON = Components()
	print buildJSON.marcDir.openDir().next()

main()

<commit_msg>added marc date parsing<commit_after>

#build core files for Presidio database

import re
import string
import os, sys
import subprocess
import xml.etree.ElementTree as ET
from xml.dom.minidom import parseString
import inspect
from pymarc import MARCReader
#may not need some of these, preserve memory
sys.setrecursionlimit(70000)

# open files, folders, read them
class fileHandlers:

	def __init__(self, directory, directoryAlt):
		self.directory = directory
		self.directoryAlt = directoryAlt #may be redundant

	'''
	#generator that returns next record in specified folder
	#turned out to be extremely buggy when abstracted
	def openDir(self, dir):
			for file in os.listdir(self.directory):
	    			if (os.path.isfile(os.path.join(self.directory, file))):
	        			yield file
	'''
	# older, slow, terrible method for iterating over file directory
	# on the wishlist to improve but for now it gets the job done
	def openDir(self, dir):
		files = []
		for item in os.listdir(dir):
        		if (os.path.isfile(os.path.join(dir, item))):
            			files.append(item)
    				return files

	#reads data within a file
	def readItem(self, count):
		file = open((os.path.join(self.directory, self.openDir
			(self.directory)[count])), 'r')
		data = file.read()
		file.close()
		return data

	#MARC data will be used as a base. If it can't find relevant data in MARC,
	#open file in XML directory instead.
	def readAlternate(self, count, extension):
		#chop off the filetype, replace it
		#note: only possible due to archive.org's impressive uniformity
		item = self.openDir(self.directory)[count]
		item = item[:-9]
		item = item + extension
		file = open((os.path.join(self.directoryAlt, self.openDir
			(self.directoryAlt)[count])), 'r')
		data = file.read()
		file.close()
		return data

#functions that construct the JSON array
class Components:

	def __init__(self, count):
		self.count = count #is an integer
		#how and I going to pass this along in recursion?

		self.files = fileHandlers('/data/MARC',
			'/home/will/MHL_download/mhl_meta_xml_files')
		self.marcNext = self.files.openDir(self.files.directory)
		self.xmlNext = self.files.openDir(self.files.directoryAlt)

		self.readMarc = self.files.readItem(count)
		self.readXml = self.files.readAlternate(count, "_meta.xml")

		self.marcRoot = ET.fromstring(self.readMarc)
		self.xmlRoot = ET.fromstring(self.readXml)

#problem: how to align, synchornize two different filetypes
	def getDate(self):
		if True or False: #leaving flexibility in, currently always true
			root = self.marcRoot
			for child in root:
				tag = child.get('tag')
				if tag == "260":
						for sub in child:
							code = sub.get('code')
							if code == "c":
								return code
		else:
			# can put alternate XML method here
			print "failed"
		pass

	def getLibrary(self):
		try:
			for contributor in xmlRoot.findall('contributor'):
				library = contributor.text
		except TypeError:
			library = ""

		# validation
		if 'library' in locals():
			pass
		else:
			library = "Other"

	def getLanguage(self):
		for language in xmlRoot.findall('language'):
			language = language.text

			if language=="English": #MHL database inconsistency
				language = "eng"

		# validation
		if 'language' in locals():
			pass
		else:
			language = ""

	def getTitle(self):
		for title in root.findall('title'):
			title = title.text
		#sets title to filename if it does not exist
		if 'title' in locals():
			title = filter(lambda x: x in string.printable, title)
			title = title.replace('"', '')
			title = title.replace("'", " ")
		else:
			title = filename

	def getSearchstring(self):
		for arlink in xmlRoot.findall('identifier-access'):
			searchstring = arlink.text
			if searchstring in locals():
				searchstring = "[" + author + "], " + title + \
					 			" (" + year + ") read"
				searchstring = searchstring.encode('ascii','ignore')
			else:
				searchstring = ""

	def getSubject(self):
		pass

	def getLocation(self):
		pass

def main():
	buildJSON = Components(0)
	print buildJSON.readMarc
	print buildJSON.getDate()

main()

<|endoftext|>


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

import parse_contactfile

class PconscContactParser(parse_contactfile.ContactfileParser):
    """ Parser for PconsC1/2/3 contact prediction files """

    _RES1 = 0
    _RES2 = 1
    _RAW_SCORE = 2

    _METHOD = "pconsc"

    def __init__(self):
        parse_contactfile.ContactfileParser.__init__(self)

    def read(self, contactfile):
        with open(contactfile, "r") as fh:
            for line in iter(fh.readline, ''):
                line = line.strip().split()

                # Define the contact in a dictionary - use parent method
                contact = self.defineContact(line,
                                             res1_idx=self._RES1,
                                             res2_idx=self._RES2,
                                             raw_score_idx=self._raw_SCORE,
                                             method=self._METHOD,
                                             file=contactfile)

                self.contacts.append(contact)
        return
##End PconscContactParser
<commit_msg>Updated variable name<commit_after>#!/usr/bin/env ccp4-python

import parse_contactfile

class PconscContactParser(parse_contactfile.ContactfileParser):
    """ Parser for PconsC1/2/3 contact prediction files """

    _RES1 = 0
    _RES2 = 1
    _RAW_SCORE = 2

    _METHOD = "pconsc"

    def __init__(self):
        parse_contactfile.ContactfileParser.__init__(self)

    def read(self, contactfile):
        with open(contactfile, "r") as fh:
            for line in iter(fh.readline, ''):
                line = line.strip().split()

                # Define the contact in a dictionary - use parent method
                contact = self.defineContact(line,
                                             res1_idx=self._RES1,
                                             res2_idx=self._RES2,
                                             raw_score_idx=self._RAW_SCORE,
                                             method=self._METHOD,
                                             file=contactfile)

                self.contacts.append(contact)
        return
##End PconscContactParser
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""MovingMNIST."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets.public_api as tfds
from tensorflow_datasets.video.moving_sequence import image_as_moving_sequence  # pylint: disable=unused-import

_OUT_RESOLUTION = (64, 64)
_SEQUENCE_LENGTH = 20
_URL = "http://www.cs.toronto.edu/~nitish/unsupervised_video/"
_CITATION = """\
@article{DBLP:journals/corr/SrivastavaMS15,
  author    = {Nitish Srivastava and
               Elman Mansimov and
               Ruslan Salakhutdinov},
  title     = {Unsupervised Learning of Video Representations using LSTMs},
  journal   = {CoRR},
  volume    = {abs/1502.04681},
  year      = {2015},
  url       = {http://arxiv.org/abs/1502.04681},
  archivePrefix = {arXiv},
  eprint    = {1502.04681},
  timestamp = {Mon, 13 Aug 2018 16:47:05 +0200},
  biburl    = {https://dblp.org/rec/bib/journals/corr/SrivastavaMS15},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}
"""
_DESCRIPTION = """\
Moving variant of MNIST database of handwritten digits. This is the
data used by the authors for reporting model performance. See
`tfds.video.moving_mnist.image_as_moving_sequence`
for generating training/validation data from the MNIST dataset.
"""

class MovingMnist(tfds.core.GeneratorBasedBuilder):
  """MovingMnist."""

  VERSION = tfds.core.Version(
      "1.0.0", "New split API (https://tensorflow.org/datasets/splits)")
  SUPPORTED_VERSIONS = [
      tfds.core.Version("0.1.0", experiments={tfds.core.Experiment.S3: False}),
  ]

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            "image_sequence": tfds.features.Video(
                shape=(_SEQUENCE_LENGTH,) + _OUT_RESOLUTION + (1,))
        }),
        homepage=_URL,
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    data_path = dl_manager.download(_URL + "mnist_test_seq.npy")

    # authors only provide test data.
    # See `tfds.video.moving_mnist.image_as_moving_sequence` for mapping
    # function to create training/validation dataset from MNIST.
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,
            num_shards=5,
            gen_kwargs=dict(data_path=data_path)),
    ]

  def _generate_examples(self, data_path):
    """Generate MovingMnist sequences.

    Args:
      data_path (str): Path to the data file

    Yields:
      20 x 64 x 64 x 1 uint8 numpy arrays
    """
    with tf.io.gfile.GFile(data_path, "rb") as fp:
      images = np.load(fp)
    images = np.transpose(images, (1, 0, 2, 3))
    images = np.expand_dims(images, axis=-1)
    for sequence in images:
      yield dict(image_sequence=sequence)
<commit_msg>Fixing moving-mnist that was broken after S3 migration.<commit_after># coding=utf-8
# Copyright 2020 The TensorFlow Datasets Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""MovingMNIST."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets.public_api as tfds
from tensorflow_datasets.video.moving_sequence import image_as_moving_sequence  # pylint: disable=unused-import

_OUT_RESOLUTION = (64, 64)
_SEQUENCE_LENGTH = 20
_URL = "http://www.cs.toronto.edu/~nitish/unsupervised_video/"
_CITATION = """\
@article{DBLP:journals/corr/SrivastavaMS15,
  author    = {Nitish Srivastava and
               Elman Mansimov and
               Ruslan Salakhutdinov},
  title     = {Unsupervised Learning of Video Representations using LSTMs},
  journal   = {CoRR},
  volume    = {abs/1502.04681},
  year      = {2015},
  url       = {http://arxiv.org/abs/1502.04681},
  archivePrefix = {arXiv},
  eprint    = {1502.04681},
  timestamp = {Mon, 13 Aug 2018 16:47:05 +0200},
  biburl    = {https://dblp.org/rec/bib/journals/corr/SrivastavaMS15},
  bibsource = {dblp computer science bibliography, https://dblp.org}
}
"""
_DESCRIPTION = """\
Moving variant of MNIST database of handwritten digits. This is the
data used by the authors for reporting model performance. See
`tfds.video.moving_mnist.image_as_moving_sequence`
for generating training/validation data from the MNIST dataset.
"""

class MovingMnist(tfds.core.GeneratorBasedBuilder):
  """MovingMnist."""

  VERSION = tfds.core.Version(
      "1.0.0", "New split API (https://tensorflow.org/datasets/splits)")
  SUPPORTED_VERSIONS = [
      tfds.core.Version("0.1.0", experiments={tfds.core.Experiment.S3: False}),
  ]

  def _info(self):
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            "image_sequence": tfds.features.Video(
                shape=(_SEQUENCE_LENGTH,) + _OUT_RESOLUTION + (1,))
        }),
        homepage=_URL,
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager):
    data_path = dl_manager.download(_URL + "mnist_test_seq.npy")

    # authors only provide test data.
    # See `tfds.video.moving_mnist.image_as_moving_sequence` for mapping
    # function to create training/validation dataset from MNIST.
    return [
        tfds.core.SplitGenerator(
            name=tfds.Split.TEST,
            gen_kwargs=dict(data_path=data_path)),
    ]

  def _generate_examples(self, data_path):
    """Generate MovingMnist sequences.

    Args:
      data_path (str): Path to the data file

    Yields:
      20 x 64 x 64 x 1 uint8 numpy arrays
    """
    with tf.io.gfile.GFile(data_path, "rb") as fp:
      images = np.load(fp)
    images = np.transpose(images, (1, 0, 2, 3))
    images = np.expand_dims(images, axis=-1)
    for i, sequence in enumerate(images):
      yield i, dict(image_sequence=sequence)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
A luigi file system client that wraps around snakebite

Originally written by Alan Brenner <alan@magnetic.com> github.com/alanbbr
"""

from luigi.contrib.hdfs import config as hdfs_config
from luigi.contrib.hdfs import error as hdfs_error
from luigi.contrib.hdfs import abstract_client as hdfs_abstract_client
from luigi import six
import luigi.contrib.target
import logging
import datetime
import os

logger = logging.getLogger('luigi-interface')

class SnakebiteHdfsClient(hdfs_abstract_client.HdfsFileSystem):
    """
    A hdfs client using snakebite. Since Snakebite has a python API, it'll be
    about 100 times faster than the hadoop cli client, which does shell out to
    a java program on each file system operation.
    """

    def __init__(self):
        super(SnakebiteHdfsClient, self).__init__()
        self._bite = None
        self.pid = -1

    @staticmethod
    def list_path(path):
        if isinstance(path, list) or isinstance(path, tuple):
            return path
        # TODO: Should this be:
        # isinstance(path, (six.text_type, six.binary_type))?
        if isinstance(path, six.string_types):
            return [path, ]
        return [str(path), ]

    def get_bite(self):
        """
        If Luigi has forked, we have a different PID, and need to reconnect.
        """
        config = hdfs_config.hdfs()
        if self.pid != os.getpid() or not self._bite:
            client_kwargs = dict(filter(
                lambda k_v: k_v[1] is not None and k_v[1] != '', six.iteritems({
                    'hadoop_version': config.client_version,
                    'effective_user': config.effective_user,
                })
            ))
            if config.snakebite_autoconfig:
                """
                This is fully backwards compatible with the vanilla Client and can be used for a non HA cluster as well.
                This client tries to read ``${HADOOP_PATH}/conf/hdfs-site.xml`` to get the address of the namenode.
                The behaviour is the same as Client.
                """
                from snakebite.client import AutoConfigClient
                self._bite = AutoConfigClient(**client_kwargs)
            else:
                from snakebite.client import Client
                self._bite = Client(config.namenode_host, config.namenode_port, **client_kwargs)
        return self._bite

    def exists(self, path):
        """
        Use snakebite.test to check file existence.

        :param path: path to test
        :type path: string
        :return: boolean, True if path exists in HDFS
        """
        try:
            return self.get_bite().test(path, exists=True)
        except Exception as err:    # IGNORE:broad-except
            raise hdfs_error.HDFSCliError("snakebite.test", -1, str(err), repr(err))

    def move(self, path, dest):
        """
        Use snakebite.rename, if available.

        :param path: source file(s)
        :type path: either a string or sequence of strings
        :param dest: destination file (single input) or directory (multiple)
        :type dest: string
        :return: list of renamed items
        """
        parts = dest.rstrip('/').split('/')
        if len(parts) > 1:
            dir_path = '/'.join(parts[0:-1])
            if not self.exists(dir_path):
                self.mkdir(dir_path, parents=True)
        return list(self.get_bite().rename(self.list_path(path), dest))

    def rename_dont_move(self, path, dest):
        """
        Use snakebite.rename_dont_move, if available.

        :param path: source path (single input)
        :type path: string
        :param dest: destination path
        :type dest: string
        :return: True if succeeded
        :raises: snakebite.errors.FileAlreadyExistsException
        """
        from snakebite.errors import FileAlreadyExistsException
        try:
            self.get_bite().rename2(path, dest, overwriteDest=False)
        except FileAlreadyExistsException:
            # Unfortunately python2 don't allow exception chaining.
            raise luigi.target.FileAlreadyExists()

    def remove(self, path, recursive=True, skip_trash=False):
        """
        Use snakebite.delete, if available.

        :param path: delete-able file(s) or directory(ies)
        :type path: either a string or a sequence of strings
        :param recursive: delete directories trees like \*nix: rm -r
        :type recursive: boolean, default is True
        :param skip_trash: do or don't move deleted items into the trash first
        :type skip_trash: boolean, default is False (use trash)
        :return: list of deleted items
        """
        return list(self.get_bite().delete(self.list_path(path), recurse=recursive))

    def chmod(self, path, permissions, recursive=False):
        """
        Use snakebite.chmod, if available.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param permissions: \*nix style permission number
        :type permissions: octal
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        if type(permissions) == str:
            permissions = int(permissions, 8)
        return list(self.get_bite().chmod(self.list_path(path),
                                          permissions, recursive))

    def chown(self, path, owner, group, recursive=False):
        """
        Use snakebite.chown/chgrp, if available.

        One of owner or group must be set. Just setting group calls chgrp.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param owner: new owner, can be blank
        :type owner: string
        :param group: new group, can be blank
        :type group: string
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        bite = self.get_bite()
        if owner:
            if group:
                return all(bite.chown(self.list_path(path), "%s:%s" % (owner, group),
                                      recurse=recursive))
            return all(bite.chown(self.list_path(path), owner, recurse=recursive))
        return list(bite.chgrp(self.list_path(path), group, recurse=recursive))

    def count(self, path):
        """
        Use snakebite.count, if available.

        :param path: directory to count the contents of
        :type path: string
        :return: dictionary with content_size, dir_count and file_count keys
        """
        try:
            res = self.get_bite().count(self.list_path(path)).next()
            dir_count = res['directoryCount']
            file_count = res['fileCount']
            content_size = res['spaceConsumed']
        except StopIteration:
            dir_count = file_count = content_size = 0
        return {'content_size': content_size, 'dir_count': dir_count,
                'file_count': file_count}

    def copy(self, path, destination):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("SnakebiteClient in luigi doesn't implement copy")

    def put(self, local_path, destination):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("Snakebite doesn't implement put")

    def get(self, path, local_destination):
        """
        Use snakebite.copyToLocal, if available.

        :param path: HDFS file
        :type path: string
        :param local_destination: path on the system running Luigi
        :type local_destination: string
        """
        return list(self.get_bite().copyToLocal(self.list_path(path),
                                                local_destination))

    def mkdir(self, path, parents=True, mode=0o755, raise_if_exists=False):
        """
        Use snakebite.mkdir, if available.

        Snakebite's mkdir method allows control over full path creation, so by
        default, tell it to build a full path to work like ``hadoop fs -mkdir``.

        :param path: HDFS path to create
        :type path: string
        :param parents: create any missing parent directories
        :type parents: boolean, default is True
        :param mode: \*nix style owner/group/other permissions
        :type mode: octal, default 0755
        """
        result = list(self.get_bite().mkdir(self.list_path(path),
                                            create_parent=parents, mode=mode))
        if raise_if_exists and "ile exists" in result[0].get('error', ''):
            raise luigi.target.FileAlreadyExists("%s exists" % (path, ))
        return result

    def listdir(self, path, ignore_directories=False, ignore_files=False,
                include_size=False, include_type=False, include_time=False,
                recursive=False):
        """
        Use snakebite.ls to get the list of items in a directory.

        :param path: the directory to list
        :type path: string
        :param ignore_directories: if True, do not yield directory entries
        :type ignore_directories: boolean, default is False
        :param ignore_files: if True, do not yield file entries
        :type ignore_files: boolean, default is False
        :param include_size: include the size in bytes of the current item
        :type include_size: boolean, default is False (do not include)
        :param include_type: include the type (d or f) of the current item
        :type include_type: boolean, default is False (do not include)
        :param include_time: include the last modification time of the current item
        :type include_time: boolean, default is False (do not include)
        :param recursive: list subdirectory contents
        :type recursive: boolean, default is False (do not recurse)
        :return: yield with a string, or if any of the include_* settings are
            true, a tuple starting with the path, and include_* items in order
        """
        bite = self.get_bite()
        for entry in bite.ls(self.list_path(path), recurse=recursive):
            if ignore_directories and entry['file_type'] == 'd':
                continue
            if ignore_files and entry['file_type'] == 'f':
                continue
            rval = [entry['path'], ]
            if include_size:
                rval.append(entry['length'])
            if include_type:
                rval.append(entry['file_type'])
            if include_time:
                rval.append(datetime.datetime.fromtimestamp(entry['modification_time'] / 1000))
            if len(rval) > 1:
                yield tuple(rval)
            else:
                yield rval[0]

    def touchz(self, path):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("SnakebiteClient in luigi doesn't implement touchz")
<commit_msg>Do not raise HDFSCliError inside snakebite client.<commit_after># -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
A luigi file system client that wraps around snakebite

Originally written by Alan Brenner <alan@magnetic.com> github.com/alanbbr
"""

from luigi.contrib.hdfs import config as hdfs_config
from luigi.contrib.hdfs import abstract_client as hdfs_abstract_client
from luigi import six
import luigi.contrib.target
import logging
import datetime
import os

logger = logging.getLogger('luigi-interface')

class SnakebiteHdfsClient(hdfs_abstract_client.HdfsFileSystem):
    """
    A hdfs client using snakebite. Since Snakebite has a python API, it'll be
    about 100 times faster than the hadoop cli client, which does shell out to
    a java program on each file system operation.
    """

    def __init__(self):
        super(SnakebiteHdfsClient, self).__init__()
        self._bite = None
        self.pid = -1

    @staticmethod
    def list_path(path):
        if isinstance(path, list) or isinstance(path, tuple):
            return path
        # TODO: Should this be:
        # isinstance(path, (six.text_type, six.binary_type))?
        if isinstance(path, six.string_types):
            return [path, ]
        return [str(path), ]

    def get_bite(self):
        """
        If Luigi has forked, we have a different PID, and need to reconnect.
        """
        config = hdfs_config.hdfs()
        if self.pid != os.getpid() or not self._bite:
            client_kwargs = dict(filter(
                lambda k_v: k_v[1] is not None and k_v[1] != '', six.iteritems({
                    'hadoop_version': config.client_version,
                    'effective_user': config.effective_user,
                })
            ))
            if config.snakebite_autoconfig:
                """
                This is fully backwards compatible with the vanilla Client and can be used for a non HA cluster as well.
                This client tries to read ``${HADOOP_PATH}/conf/hdfs-site.xml`` to get the address of the namenode.
                The behaviour is the same as Client.
                """
                from snakebite.client import AutoConfigClient
                self._bite = AutoConfigClient(**client_kwargs)
            else:
                from snakebite.client import Client
                self._bite = Client(config.namenode_host, config.namenode_port, **client_kwargs)
        return self._bite

    def exists(self, path):
        """
        Use snakebite.test to check file existence.

        :param path: path to test
        :type path: string
        :return: boolean, True if path exists in HDFS
        """
        return self.get_bite().test(path, exists=True)

    def move(self, path, dest):
        """
        Use snakebite.rename, if available.

        :param path: source file(s)
        :type path: either a string or sequence of strings
        :param dest: destination file (single input) or directory (multiple)
        :type dest: string
        :return: list of renamed items
        """
        parts = dest.rstrip('/').split('/')
        if len(parts) > 1:
            dir_path = '/'.join(parts[0:-1])
            if not self.exists(dir_path):
                self.mkdir(dir_path, parents=True)
        return list(self.get_bite().rename(self.list_path(path), dest))

    def rename_dont_move(self, path, dest):
        """
        Use snakebite.rename_dont_move, if available.

        :param path: source path (single input)
        :type path: string
        :param dest: destination path
        :type dest: string
        :return: True if succeeded
        :raises: snakebite.errors.FileAlreadyExistsException
        """
        from snakebite.errors import FileAlreadyExistsException
        try:
            self.get_bite().rename2(path, dest, overwriteDest=False)
        except FileAlreadyExistsException:
            # Unfortunately python2 don't allow exception chaining.
            raise luigi.target.FileAlreadyExists()

    def remove(self, path, recursive=True, skip_trash=False):
        """
        Use snakebite.delete, if available.

        :param path: delete-able file(s) or directory(ies)
        :type path: either a string or a sequence of strings
        :param recursive: delete directories trees like \*nix: rm -r
        :type recursive: boolean, default is True
        :param skip_trash: do or don't move deleted items into the trash first
        :type skip_trash: boolean, default is False (use trash)
        :return: list of deleted items
        """
        return list(self.get_bite().delete(self.list_path(path), recurse=recursive))

    def chmod(self, path, permissions, recursive=False):
        """
        Use snakebite.chmod, if available.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param permissions: \*nix style permission number
        :type permissions: octal
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        if type(permissions) == str:
            permissions = int(permissions, 8)
        return list(self.get_bite().chmod(self.list_path(path),
                                          permissions, recursive))

    def chown(self, path, owner, group, recursive=False):
        """
        Use snakebite.chown/chgrp, if available.

        One of owner or group must be set. Just setting group calls chgrp.

        :param path: update-able file(s)
        :type path: either a string or sequence of strings
        :param owner: new owner, can be blank
        :type owner: string
        :param group: new group, can be blank
        :type group: string
        :param recursive: change just listed entry(ies) or all in directories
        :type recursive: boolean, default is False
        :return: list of all changed items
        """
        bite = self.get_bite()
        if owner:
            if group:
                return all(bite.chown(self.list_path(path), "%s:%s" % (owner, group),
                                      recurse=recursive))
            return all(bite.chown(self.list_path(path), owner, recurse=recursive))
        return list(bite.chgrp(self.list_path(path), group, recurse=recursive))

    def count(self, path):
        """
        Use snakebite.count, if available.

        :param path: directory to count the contents of
        :type path: string
        :return: dictionary with content_size, dir_count and file_count keys
        """
        try:
            res = self.get_bite().count(self.list_path(path)).next()
            dir_count = res['directoryCount']
            file_count = res['fileCount']
            content_size = res['spaceConsumed']
        except StopIteration:
            dir_count = file_count = content_size = 0
        return {'content_size': content_size, 'dir_count': dir_count,
                'file_count': file_count}

    def copy(self, path, destination):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("SnakebiteClient in luigi doesn't implement copy")

    def put(self, local_path, destination):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("Snakebite doesn't implement put")

    def get(self, path, local_destination):
        """
        Use snakebite.copyToLocal, if available.

        :param path: HDFS file
        :type path: string
        :param local_destination: path on the system running Luigi
        :type local_destination: string
        """
        return list(self.get_bite().copyToLocal(self.list_path(path),
                                                local_destination))

    def mkdir(self, path, parents=True, mode=0o755, raise_if_exists=False):
        """
        Use snakebite.mkdir, if available.

        Snakebite's mkdir method allows control over full path creation, so by
        default, tell it to build a full path to work like ``hadoop fs -mkdir``.

        :param path: HDFS path to create
        :type path: string
        :param parents: create any missing parent directories
        :type parents: boolean, default is True
        :param mode: \*nix style owner/group/other permissions
        :type mode: octal, default 0755
        """
        result = list(self.get_bite().mkdir(self.list_path(path),
                                            create_parent=parents, mode=mode))
        if raise_if_exists and "ile exists" in result[0].get('error', ''):
            raise luigi.target.FileAlreadyExists("%s exists" % (path, ))
        return result

    def listdir(self, path, ignore_directories=False, ignore_files=False,
                include_size=False, include_type=False, include_time=False,
                recursive=False):
        """
        Use snakebite.ls to get the list of items in a directory.

        :param path: the directory to list
        :type path: string
        :param ignore_directories: if True, do not yield directory entries
        :type ignore_directories: boolean, default is False
        :param ignore_files: if True, do not yield file entries
        :type ignore_files: boolean, default is False
        :param include_size: include the size in bytes of the current item
        :type include_size: boolean, default is False (do not include)
        :param include_type: include the type (d or f) of the current item
        :type include_type: boolean, default is False (do not include)
        :param include_time: include the last modification time of the current item
        :type include_time: boolean, default is False (do not include)
        :param recursive: list subdirectory contents
        :type recursive: boolean, default is False (do not recurse)
        :return: yield with a string, or if any of the include_* settings are
            true, a tuple starting with the path, and include_* items in order
        """
        bite = self.get_bite()
        for entry in bite.ls(self.list_path(path), recurse=recursive):
            if ignore_directories and entry['file_type'] == 'd':
                continue
            if ignore_files and entry['file_type'] == 'f':
                continue
            rval = [entry['path'], ]
            if include_size:
                rval.append(entry['length'])
            if include_type:
                rval.append(entry['file_type'])
            if include_time:
                rval.append(datetime.datetime.fromtimestamp(entry['modification_time'] / 1000))
            if len(rval) > 1:
                yield tuple(rval)
            else:
                yield rval[0]

    def touchz(self, path):
        """
        Raise a NotImplementedError exception.
        """
        raise NotImplementedError("SnakebiteClient in luigi doesn't implement touchz")
<|endoftext|>


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

"""
Copyright [2009-2018] EMBL-European Bioinformatics Institute
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
     http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import luigi

from .genome_mapping_tasks import GetFasta
from .genome_mapping_tasks import CleanSplitFasta
from .genome_mapping_tasks import GetChromosomes
from .genome_mapping_tasks import BlatJob
from .genome_mapping_tasks import ParsePslOutput

from .pgload_exact_matches import GenomeMappingPGLoadExactMatches
from .pgload_inexact_matches import GenomeMappingPGLoadInexactMatches

def get_taxids_for_genome_mapping():
    """
    Get taxids for genomes that are used for mapping.
    """
    return [
        4932, # S. cerevisiae
        # 9606, # human
        # 10090, # mouse
        # 10116, # rat
    ]

class SpeciesFastaExportWrapper(luigi.WrapperTask):
    """
    A wrapper task to export fasta files for all species that will be mapped
    to the reference genomes using blat.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield GetFasta(taxid=taxid)

class SpeciesFastaCleanSplitWrapper(luigi.WrapperTask):
    """
    A wrapper task to keep only sequences of certain length and split fasta
    files in chunks.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield CleanSplitFasta(taxid=taxid)

class GetChromosomeFastaWrapper(luigi.WrapperTask):
    """
    A wrapper task for getting a list of all chromosome fasta files
    that are used in parallel blat searches.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield GetChromosomes(taxid=taxid)

class BlatJobsWrapper(luigi.WrapperTask):
    """
    A wrapper task for running blat searches of all split RNAcentral fasta files
    against all chromosomes within the same species.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            for chunk in CleanSplitFasta(taxid=taxid).output():
                for chromosome in GetChromosomes(taxid=taxid).output():
                    yield BlatJob(fasta_input=chunk.path,
                                  chromosome=chromosome.path,
                                  taxid=taxid)

class PGLoadGenomeMappingWrapper(luigi.WrapperTask):
    """
    A wrapper task for loading parsed blat output into the database.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield [
                GenomeMappingPGLoadExactMatches(taxid=taxid),
                GenomeMappingPGLoadInexactMatches(taxid=taxid),
            ]

class GenomeMappingPipelineWrapper(luigi.WrapperTask):
    """
    A wrapper task for the entire genome mapping pipeline.
    """
    def requires(self):
        yield GetChromosomeFastaWrapper()
        yield SpeciesFastaCleanSplitWrapper()
        yield BlatJobsWrapper()
        yield GenomeMappingPGLoadExactMatches()
        yield GenomeMappingPGLoadInexactMatches()
<commit_msg>Add a wrapper task for parsing blat output<commit_after># -*- coding: utf-8 -*-

"""
Copyright [2009-2018] EMBL-European Bioinformatics Institute
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
     http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import luigi

from .genome_mapping_tasks import GetFasta
from .genome_mapping_tasks import CleanSplitFasta
from .genome_mapping_tasks import GetChromosomes
from .genome_mapping_tasks import BlatJob
from .genome_mapping_tasks import ParsePslOutput

from .pgload_exact_matches import GenomeMappingPGLoadExactMatches
from .pgload_inexact_matches import GenomeMappingPGLoadInexactMatches

def get_taxids_for_genome_mapping():
    """
    Get taxids for genomes that are used for mapping.
    """
    return [
        4932, # S. cerevisiae
        # 9606, # human
        # 10090, # mouse
        # 10116, # rat
    ]

class SpeciesFastaExportWrapper(luigi.WrapperTask):
    """
    A wrapper task to export fasta files for all species that will be mapped
    to the reference genomes using blat.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield GetFasta(taxid=taxid)

class SpeciesFastaCleanSplitWrapper(luigi.WrapperTask):
    """
    A wrapper task to keep only sequences of certain length and split fasta
    files in chunks.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield CleanSplitFasta(taxid=taxid)

class GetChromosomeFastaWrapper(luigi.WrapperTask):
    """
    A wrapper task for getting a list of all chromosome fasta files
    that are used in parallel blat searches.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield GetChromosomes(taxid=taxid)

class BlatJobsWrapper(luigi.WrapperTask):
    """
    A wrapper task for running blat searches of all split RNAcentral fasta files
    against all chromosomes within the same species.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            for chunk in CleanSplitFasta(taxid=taxid).output():
                for chromosome in GetChromosomes(taxid=taxid).output():
                    yield BlatJob(fasta_input=chunk.path,
                                  chromosome=chromosome.path,
                                  taxid=taxid)

class ParsePslOutputWrapper(luigi.WrapperTask):
    """
    A wrapper task for parsing all blat output into tsv files that can be
    loaded into the database.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield ParsePslOutput(taxid=taxid)

class PGLoadGenomeMappingWrapper(luigi.WrapperTask):
    """
    A wrapper task for loading parsed blat output into the database.
    """
    def requires(self):
        for taxid in get_taxids_for_genome_mapping():
            yield [
                GenomeMappingPGLoadExactMatches(taxid=taxid),
                GenomeMappingPGLoadInexactMatches(taxid=taxid),
            ]

class GenomeMappingPipelineWrapper(luigi.WrapperTask):
    """
    A wrapper task for the entire genome mapping pipeline.
    """
    def requires(self):
        yield GetChromosomeFastaWrapper()
        yield SpeciesFastaCleanSplitWrapper()
        yield BlatJobsWrapper()
        yield ParsePslOutputWrapper()
        yield GenomeMappingPGLoadExactMatches()
        yield GenomeMappingPGLoadInexactMatches()
<|endoftext|>