<commit_before>__author__ = "AstroPrint Product Team <product@astroprint.com>"
__license__ = "GNU Affero General Public License http://www.gnu.org/licenses/agpl.html"
__copyright__ = "Copyright (C) 2018 3DaGoGo, Inc - Released under terms of the AGPLv3 License"

import logging
import os
import threading
import sarge
import pyudev
import fnmatch

from glob import glob

from astroprint.printer.manager import printerManager

from octoprint.events import eventManager, Events
from octoprint.settings import settings

ROOT_MOUNT_POINT = '/media/astrobox'

#
# Thread to get some plugged usb drive
#
class ExternalDriveManager(threading.Thread):
	def __init__(self):
		super(ExternalDriveManager, self).__init__()
		self.daemon = True
		self.stopThread = False
		self.monitor = None
		self._mountPoints = None
		self._logger = logging.getLogger(__name__)
		self._eventManager = eventManager()
		self._logger.info('Starting Linux ExternalDriveManager')

	def run(self):
		context = pyudev.Context()

		#Check the status of the system right now
		# 1. Check if we have any drives connected
		for d in context.list_devices(subsystem='block', DEVTYPE='disk'):
			if d.attributes.asbool('removable'):
				#it's removable media, let's find the partitions
				partitions = list(context.list_devices(subsystem='block', DEVTYPE='partition', parent=d))
				for p in partitions:
					# we only analyze the first one, ignore other partitions
					#p = partitions[0]

					#check if the partition is mounted
					mountPoint = self._findMountPoint(p.device_node)
					if mountPoint:
						self._logger.info('Found mounted removable drive (%s) at %s' % (p.device_node, mountPoint))
					else:
						self._logger.info('Mounting inserted removable drive (%s)' % p.device_node)
						self._mountPartition(p.device_node, self._getDeviceMountDirectory(p))

		# 2. Check if there are left over directories with no drives mounted
		for f in glob('%s/*/*' % ROOT_MOUNT_POINT):
			if os.path.isdir(f) and not os.listdir(f):
				#empty directory found, let's see if it's mounted
				if not self._isMounted(f):
					#Not mounted and empty, delete directory.
					os.rmdir(f) #main dir
					os.rmdir('/'.join(f.split('/')[:-1])) #uuid dir

		self._mountPoints = None # We reset here since we don't need it anymore, it will be re-created on shutdown

		# Start listening for events
		self.monitor = pyudev.Monitor.from_netlink(context)
		self.monitor.filter_by(subsystem='block')

		for device in iter(self.monitor.poll, None):

			if self.stopThread:
				self.monitor.stop()
				return

			if device.device_type == 'partition':
				if device.action == 'add':
					devName = device.device_node
					self._logger.info('%s pluged in' % devName)
					mountPath = self._getDeviceMountDirectory(device)
					if self._mountPartition(devName, mountPath):
						self._eventManager.fire( Events.EXTERNAL_DRIVE_MOUNTED, {
							"mount_path": mountPath,
							"device_node": devName
						})

				if device.action == 'remove':
					devName = device.device_node
					mountPath = self._getDeviceMountDirectory(device)
					self._logger.info('%s removed' % devName)
					if self._umountPartition(self._getDeviceMountDirectory(device)):
						self._eventManager.fire( Events.EXTERNAL_DRIVE_PHISICALLY_REMOVED, {
							"device_node": devName,
							"mount_path": mountPath,
						})

	def _getDeviceMountDirectory(self, device):
		name = device.get('ID_FS_LABEL', 'NO NAME')
		uuid = device.get('ID_FS_UUID')
		if name and uuid:
			return os.path.join(ROOT_MOUNT_POINT, uuid, name)
		else:
			return None

	def _loadMountPoints(self):
		self._mountPoints = {}
		with open('/proc/mounts', 'rt') as f:
			for line in f:
				line = line.strip()
				parts = line.split(' ')

				self._mountPoints[parts[0]] = parts[1]

	def _findMountPoint(self, devPath):
		if self._mountPoints is None:
			self._loadMountPoints()

		return self._mountPoints.get(devPath)

	def _isMounted(self, dirPath):
		if self._mountPoints is None:
			self._loadMountPoints()

		for dev, path in self._mountPoints.iteritems():
			if path == dirPath:
				return True

		return False

	def _mountPartition(self, partition, directory):
		if directory:
			try:
				if not os.path.exists(directory):
					os.makedirs(directory)

				p = sarge.run("mount %s '%s'" % (partition, directory), stderr=sarge.Capture())
				if p.returncode != 0:
					returncode = p.returncode
					stderr_text = p.stderr.text
					self._logger.warn("Partition mount failed with return code %i: %s" % (returncode, stderr_text))
					return False

				else:
					self._logger.info("Partition %s mounted on %s" % (partition, directory))
					return True

			except Exception, e:
				self._logger.warn("Mount failed: %s" % e)
				return False

		else:
			return False

	def _umountPartition(self, directory):
		if directory:
			try:
				if os.path.exists(directory):
					p = sarge.run("umount '%s'" % directory, stderr=sarge.Capture())
					if p.returncode != 0:
						returncode = p.returncode
						stderr_text = p.stderr.text
						self._logger.warn("Partition umount failed with return code %i: %s" % (returncode, stderr_text))
						return False

					else:
						os.rmdir(directory)
						os.rmdir('/'.join(directory.split('/')[:-1])) #uuid dir
						self._logger.info("Partition umounted from %s" % directory)
						return True

				else:
					return True

			except Exception, e:
				self._logger.warn("umount failed: %s" % e)
				return False

		else:
			return False

	def shutdown(self):
		self._logger.info('Shutting Down Linux ExternalDriveManager')

		#Unmount mounted drives
		for d in glob('%s/*/*' % ROOT_MOUNT_POINT):
			if self._isMounted(d):
				self._umountPartition(d)

		self.stopThread = True

	def _cleanFileLocation(self, location):
		locationParsed = location.replace('//','/')
		return locationParsed

	def eject(self, mountPath):

		retries = 5
		timeout = 3

		ejected = False

		for i in (0,retries):
			if self._umountPartition(mountPath):
				ejected = True
			else:
				sleep(timeout)

		if ejected:
			self._eventManager.fire( Events.EXTERNAL_DRIVE_EJECTED, {
				"mount_path": mountPath
			})
			return {'result': True}

		return {
			'result': False,
			'error': "Unable to eject"
		}

	def copy(self, src, dst, progressCb, observerId):
			blksize = 1048576 # 1MiB
			try:
					s = open(src, 'rb')
					d = open(dst, 'wb')
					sizeWritten = 0
			except (KeyboardInterrupt, Exception) as e:
					if 's' in locals():
							s.close()
					if 'd' in locals():
							d.close()
					raise
			try:
					total = float(os.stat(src).st_size)

					while sizeWritten <= total:
							buf = s.read(blksize)
							bytes_written = d.write(buf)

							progressCb(int((os.stat(dst).st_size / total)*100),dst,observerId)

							sizeWritten += blksize

							if blksize > len(buf) or bytes_written == 0:
									printerManager().fileManager._metadataAnalyzer.addFileToQueue(dst)
									progressCb(100,dst,observerId)
									break

			except (KeyboardInterrupt, Exception) as e:
					s.close()
					d.close()
					raise
			else:
					progressCb(100,dst,observerId)
					s.close()
					d.close()

	def localFileExists(self, filename):
		try:
				s = open(self._cleanFileLocation(self.getBaseFolder('uploads') + '/' + filename), 'rb')
		except Exception as e:
				if 's' in locals():
					s.close()

				return False

		s.close()

		return True

	def _progressCb(self, progress,file,observerId):
		self._eventManager.fire(
			Events.COPY_TO_HOME_PROGRESS, {
				"type": "progress",
				"file": file,
				"observerId": observerId,
				"progress": progress
			}
		)

	def copyFileToLocal(self, origin, destination, observerId):
		try:
			_origin = self._cleanFileLocation(origin)

			self.copy(_origin,self._cleanFileLocation(destination)+'/'+origin.split('/')[-1:][0],self._progressCb,observerId)

			return True

		except Exception as e:
			self._logger.error("copy print file to local folder failed", exc_info = True)

			return False

	def getRemovableDrives(self):
		return self.getDirContents('%s/*/*' % ROOT_MOUNT_POINT, 'usb')

	def getFileBrowsingExtensions(self):
		return printerManager().fileManager.SUPPORTED_EXTENSIONS

	def getFolderContents(self, folder):
		try:
			return self.getDirContents(self._cleanFileLocation(folder))

		except Exception as e:
			self._logger.error("exploration folders can not be obtained", exc_info = True)
			return None

	def getBaseFolder(self, key):
		return self._cleanFileLocation(settings().getBaseFolder(key))

	def getDirContents(self, globPattern, icon='folder', extensions=None):
		if extensions is None:
			extensions = self.getFileBrowsingExtensions()

		files = []

		for item in glob(globPattern):
			if os.path.isdir(item):
				files.append({
					'name': item,
					'icon': icon})

			else:
				for ext in extensions:
					if fnmatch.fnmatch(item.lower(), '*.' + ext):
						files.append({
							'name': item,
							'icon': ext
						})
						break

		return files
<commit_msg>missing import when eject needs to be retried<commit_after>__author__ = "AstroPrint Product Team <product@astroprint.com>"
__license__ = "GNU Affero General Public License http://www.gnu.org/licenses/agpl.html"
__copyright__ = "Copyright (C) 2018 3DaGoGo, Inc - Released under terms of the AGPLv3 License"

import logging
import os
import threading
import sarge
import pyudev
import fnmatch
import time

from glob import glob

from astroprint.printer.manager import printerManager

from octoprint.events import eventManager, Events
from octoprint.settings import settings

ROOT_MOUNT_POINT = '/media/astrobox'

#
# Thread to get some plugged usb drive
#
class ExternalDriveManager(threading.Thread):
	def __init__(self):
		super(ExternalDriveManager, self).__init__()
		self.daemon = True
		self.stopThread = False
		self.monitor = None
		self._mountPoints = None
		self._logger = logging.getLogger(__name__)
		self._eventManager = eventManager()
		self._logger.info('Starting Linux ExternalDriveManager')

	def run(self):
		context = pyudev.Context()

		#Check the status of the system right now
		# 1. Check if we have any drives connected
		for d in context.list_devices(subsystem='block', DEVTYPE='disk'):
			if d.attributes.asbool('removable'):
				#it's removable media, let's find the partitions
				partitions = list(context.list_devices(subsystem='block', DEVTYPE='partition', parent=d))
				for p in partitions:
					# we only analyze the first one, ignore other partitions
					#p = partitions[0]

					#check if the partition is mounted
					mountPoint = self._findMountPoint(p.device_node)
					if mountPoint:
						self._logger.info('Found mounted removable drive (%s) at %s' % (p.device_node, mountPoint))
					else:
						self._logger.info('Mounting inserted removable drive (%s)' % p.device_node)
						self._mountPartition(p.device_node, self._getDeviceMountDirectory(p))

		# 2. Check if there are left over directories with no drives mounted
		for f in glob('%s/*/*' % ROOT_MOUNT_POINT):
			if os.path.isdir(f) and not os.listdir(f):
				#empty directory found, let's see if it's mounted
				if not self._isMounted(f):
					#Not mounted and empty, delete directory.
					os.rmdir(f) #main dir
					os.rmdir('/'.join(f.split('/')[:-1])) #uuid dir

		self._mountPoints = None # We reset here since we don't need it anymore, it will be re-created on shutdown

		# Start listening for events
		self.monitor = pyudev.Monitor.from_netlink(context)
		self.monitor.filter_by(subsystem='block')

		for device in iter(self.monitor.poll, None):

			if self.stopThread:
				self.monitor.stop()
				return

			if device.device_type == 'partition':
				if device.action == 'add':
					devName = device.device_node
					self._logger.info('%s pluged in' % devName)
					mountPath = self._getDeviceMountDirectory(device)
					if self._mountPartition(devName, mountPath):
						self._eventManager.fire( Events.EXTERNAL_DRIVE_MOUNTED, {
							"mount_path": mountPath,
							"device_node": devName
						})

				if device.action == 'remove':
					devName = device.device_node
					mountPath = self._getDeviceMountDirectory(device)
					self._logger.info('%s removed' % devName)
					if self._umountPartition(self._getDeviceMountDirectory(device)):
						self._eventManager.fire( Events.EXTERNAL_DRIVE_PHISICALLY_REMOVED, {
							"device_node": devName,
							"mount_path": mountPath,
						})

	def _getDeviceMountDirectory(self, device):
		name = device.get('ID_FS_LABEL', 'NO NAME')
		uuid = device.get('ID_FS_UUID')
		if name and uuid:
			return os.path.join(ROOT_MOUNT_POINT, uuid, name)
		else:
			return None

	def _loadMountPoints(self):
		self._mountPoints = {}
		with open('/proc/mounts', 'rt') as f:
			for line in f:
				line = line.strip()
				parts = line.split(' ')

				self._mountPoints[parts[0]] = parts[1]

	def _findMountPoint(self, devPath):
		if self._mountPoints is None:
			self._loadMountPoints()

		return self._mountPoints.get(devPath)

	def _isMounted(self, dirPath):
		if self._mountPoints is None:
			self._loadMountPoints()

		for dev, path in self._mountPoints.iteritems():
			if path == dirPath:
				return True

		return False

	def _mountPartition(self, partition, directory):
		if directory:
			try:
				if not os.path.exists(directory):
					os.makedirs(directory)

				p = sarge.run("mount %s '%s'" % (partition, directory), stderr=sarge.Capture())
				if p.returncode != 0:
					returncode = p.returncode
					stderr_text = p.stderr.text
					self._logger.warn("Partition mount failed with return code %i: %s" % (returncode, stderr_text))
					return False

				else:
					self._logger.info("Partition %s mounted on %s" % (partition, directory))
					return True

			except Exception, e:
				self._logger.warn("Mount failed: %s" % e)
				return False

		else:
			return False

	def _umountPartition(self, directory):
		if directory:
			try:
				if os.path.exists(directory):
					p = sarge.run("umount '%s'" % directory, stderr=sarge.Capture())
					if p.returncode != 0:
						returncode = p.returncode
						stderr_text = p.stderr.text
						self._logger.warn("Partition umount failed with return code %i: %s" % (returncode, stderr_text))
						return False

					else:
						os.rmdir(directory)
						os.rmdir('/'.join(directory.split('/')[:-1])) #uuid dir
						self._logger.info("Partition umounted from %s" % directory)
						return True

				else:
					return True

			except Exception, e:
				self._logger.warn("umount failed: %s" % e)
				return False

		else:
			return False

	def shutdown(self):
		self._logger.info('Shutting Down Linux ExternalDriveManager')

		#Unmount mounted drives
		for d in glob('%s/*/*' % ROOT_MOUNT_POINT):
			if self._isMounted(d):
				self._umountPartition(d)

		self.stopThread = True

	def _cleanFileLocation(self, location):
		locationParsed = location.replace('//','/')
		return locationParsed

	def eject(self, mountPath):

		retries = 5
		timeout = 3

		ejected = False

		for i in (0,retries):
			if self._umountPartition(mountPath):
				ejected = True
			else:
				time.sleep(timeout)

		if ejected:
			self._eventManager.fire( Events.EXTERNAL_DRIVE_EJECTED, {
				"mount_path": mountPath
			})
			return {'result': True}

		return {
			'result': False,
			'error': "Unable to eject"
		}

	def copy(self, src, dst, progressCb, observerId):
			blksize = 1048576 # 1MiB
			try:
					s = open(src, 'rb')
					d = open(dst, 'wb')
					sizeWritten = 0
			except (KeyboardInterrupt, Exception) as e:
					if 's' in locals():
							s.close()
					if 'd' in locals():
							d.close()
					raise
			try:
					total = float(os.stat(src).st_size)

					while sizeWritten <= total:
							buf = s.read(blksize)
							bytes_written = d.write(buf)

							progressCb(int((os.stat(dst).st_size / total)*100),dst,observerId)

							sizeWritten += blksize

							if blksize > len(buf) or bytes_written == 0:
									printerManager().fileManager._metadataAnalyzer.addFileToQueue(dst)
									progressCb(100,dst,observerId)
									break

			except (KeyboardInterrupt, Exception) as e:
					s.close()
					d.close()
					raise
			else:
					progressCb(100,dst,observerId)
					s.close()
					d.close()

	def localFileExists(self, filename):
		try:
				s = open(self._cleanFileLocation(self.getBaseFolder('uploads') + '/' + filename), 'rb')
		except Exception as e:
				if 's' in locals():
					s.close()

				return False

		s.close()

		return True

	def _progressCb(self, progress,file,observerId):
		self._eventManager.fire(
			Events.COPY_TO_HOME_PROGRESS, {
				"type": "progress",
				"file": file,
				"observerId": observerId,
				"progress": progress
			}
		)

	def copyFileToLocal(self, origin, destination, observerId):
		try:
			_origin = self._cleanFileLocation(origin)

			self.copy(_origin,self._cleanFileLocation(destination)+'/'+origin.split('/')[-1:][0],self._progressCb,observerId)

			return True

		except Exception as e:
			self._logger.error("copy print file to local folder failed", exc_info = True)

			return False

	def getRemovableDrives(self):
		return self.getDirContents('%s/*/*' % ROOT_MOUNT_POINT, 'usb')

	def getFileBrowsingExtensions(self):
		return printerManager().fileManager.SUPPORTED_EXTENSIONS

	def getFolderContents(self, folder):
		try:
			return self.getDirContents(self._cleanFileLocation(folder))

		except Exception as e:
			self._logger.error("exploration folders can not be obtained", exc_info = True)
			return None

	def getBaseFolder(self, key):
		return self._cleanFileLocation(settings().getBaseFolder(key))

	def getDirContents(self, globPattern, icon='folder', extensions=None):
		if extensions is None:
			extensions = self.getFileBrowsingExtensions()

		files = []

		for item in glob(globPattern):
			if os.path.isdir(item):
				files.append({
					'name': item,
					'icon': icon})

			else:
				for ext in extensions:
					if fnmatch.fnmatch(item.lower(), '*.' + ext):
						files.append({
							'name': item,
							'icon': ext
						})
						break

		return files
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Remove horizontal stripes from ACS WFC post-SM4 data.

Examples
--------

In Python without TEAL:

>>> from acstools import acs_destripe
>>> acs_destripe.clean('uncorrected_flt.fits', 'csck', clobber=False,
...                    maxiter=15, sigrej=2.0)

In Python with TEAL:

>>> from acstools import acs_destripe
>>> from stsci.tools import teal
>>> teal.teal('acs_destripe')

In Pyraf::

    --> import acstools
    --> teal acs_destripe

From command line::

    % ./acs_destripe [-h][-c] input output [maxiter # [sigrej #]]

"""
from __future__ import print_function

# STDLIB
import logging
import os
import sys

# THIRD-PARTY
import numpy as np
from astropy.io import fits

# STSCI
from stsci.tools import parseinput,teal

__taskname__ = 'acs_destripe'
__version__ = '0.2.5'
__vdate__ = '23-Apr-2013'
__author__ = 'Norman Grogin, STScI, March 2012.'

MJD_SM4 = 54967

logging.basicConfig()
LOG = logging.getLogger(__taskname__)
LOG.setLevel(logging.INFO)

class StripeArray(object):
    """Class to handle data array to be destriped."""

    def __init__(self, image):
        self.hdulist = fits.open(image)
        self.ampstring = self.hdulist[0].header['CCDAMP']
        self.flatcorr = self.hdulist[0].header['FLATCORR']
        self.configure_arrays()

    def configure_arrays(self):
        """Get the SCI and ERR data."""
        self.science = self.hdulist['sci',1].data
        self.err = self.hdulist['err',1].data
        if (self.ampstring == 'ABCD'):
            self.science = np.concatenate(
                (self.science, self.hdulist['sci',2].data[::-1,:]), axis=1)
            self.err = np.concatenate(
                (self.err, self.hdulist['err',2].data[::-1,:]), axis=1)
        self.ingest_flatfield()

    def ingest_flatfield(self):
        """Process flatfield."""

        flatfile = self.hdulist[0].header['PFLTFILE']

        # if BIAS or DARK, set flatfield to unity
        if flatfile == 'N/A':
            self.invflat = np.ones(self.science.shape)
            return

        hduflat = self.resolve_flatname(flatfile)

        if (self.ampstring == 'ABCD'):
            self.invflat = np.concatenate(
                (1 / hduflat['sci',1].data,
                 1 / hduflat['sci',2].data[::-1,:]), axis=1)
        else:
            # complex algorithm to determine proper subarray of flatfield to use

            # which amp?
            if (self.ampstring == 'A' or self.ampstring == 'B' or
                    self.ampstring == 'AB'):
                self.invflat = 1 / hduflat['sci',2].data
            else:
                self.invflat = 1 / hduflat['sci',1].data

            # now, which section?
            sizaxis1 = self.hdulist[1].header['SIZAXIS1']
            sizaxis2 = self.hdulist[1].header['SIZAXIS2']
            centera1 = self.hdulist[1].header['CENTERA1']
            centera2 = self.hdulist[1].header['CENTERA2']

            flataxis1 = hduflat[1].header['NAXIS1']
            flataxis2 = hduflat[1].header['NAXIS2']

            # configure the offset appropriate to left- or right-side of CCD
            if (self.ampstring[0] == 'A' or self.ampstring[0] == 'C'):
                xdelta = 13
            else:
                xdelta = 35

            if sizaxis1 == flataxis1:
                xlo = 0
                xhi = sizaxis1
            else:
                xlo = centera1 - xdelta - sizaxis1 / 2 - 1
                xhi = centera1 - xdelta + sizaxis1 / 2 - 1

            if sizaxis2 == flataxis2:
                ylo = 0
                yhi = sizaxis2
            else:
                ylo = centera2 - sizaxis2 / 2 - 1
                yhi = centera2 + sizaxis2 / 2 - 1

            self.invflat = self.invflat[ylo:yhi,xlo:xhi]

        # apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science / self.invflat
            self.err = self.err / self.invflat

    def resolve_flatname(self, flatfile):
        """Resolve the filename into an absolute pathname (if necessary)."""
        flatparts = flatfile.partition('$')
        if flatparts[1] == '$':
            flatdir = os.getenv(flatparts[0])
            if flatdir is None:
                errstr = 'Environment variable ' + flatparts[0] + ' undefined!'
                raise ValueError(errstr)
            flatfile = flatdir + flatparts[2]
        if not os.path.exists(flatfile):
            errstr = 'Flat-field file ' + flatfile + ' could not be found!'
            raise IOError, errstr

        return(fits.open(flatfile))

    def write_corrected(self, output, clobber=False):
        """Write out the destriped data."""

        # un-apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science * self.invflat
            self.err = self.err * self.invflat

        # reverse the amp merge
        if (self.ampstring == 'ABCD'):
            tmp_1, tmp_2 = np.split(self.science, 2, axis=1)
            self.hdulist['sci',1].data = tmp_1.copy()
            self.hdulist['sci',2].data = tmp_2[::-1,:].copy()

            tmp_1, tmp_2 = np.split(self.err, 2, axis=1)
            self.hdulist['err',1].data = tmp_1.copy()
            self.hdulist['err',2].data = tmp_2[::-1,:].copy()

        # Write the output
        self.hdulist.writeto(output, clobber=clobber)

def clean(input, suffix, maxiter=15, sigrej=2.0, clobber=False):
    """
    Remove horizontal stripes from ACS WFC post-SM4 data.

    .. note::

        Input data must be an ACS/WFC FLT image, with 2 SCI extensions.
        Does not work on RAW image.

        Uses the flatfield specified by the image header keyword PFLTFILE.
        If keyword value is 'N/A', as is the case with biases and darks,
        then unity flatfield is used.

    Parameters
    ----------
    input : str or list of str
        Input filenames in one of these formats:

            * a Python list of filenames
            * a partial filename with wildcards ('\*flt.fits')
            * filename of an ASN table ('j12345670_asn.fits')
            * an at-file (``@input``)

    suffix : str
        The string to use to add to each input file name to
        indicate an output product. This string will be appended
        to the _flt suffix in each input filename to create the
        new output filename. For example, setting `suffix='csck'`
        will create '\*_flt_csck.fits' images.

    maxiter : int
        This parameter controls the maximum number of iterations
        to perform when computing the statistics used to compute the
        row-by-row corrections.

    sigrej : float
        This parameters sets the sigma level for the rejection applied
        during each iteration of statistics computations for the
        row-by-row corrections.

    clobber : bool
        Specify whether or not to 'clobber' (delete then replace)
        previously generated products with the same names.

    """
    flist, alist = parseinput.parseinput(input)

    for image in flist:
        # Skip processing pre-SM4 images
        if (fits.getval(image, 'EXPSTART') <= MJD_SM4):
            LOG.warn(image + ' is pre-SM4. Skipping...'%image)
            continue

        # Skip processing non-DARKCORR-ed non-BIAS images
        if (fits.getval(image, 'IMAGETYP') != 'BIAS' and
                fits.getval(image,'DARKCORR') != 'COMPLETE'):
            LOG.warn(image + ' is not BIAS and does not have '
                         'DARKCORR=COMPLETE. Skipping...')
            continue

        # Data must be in ELECTRONS
        if (fits.getval(image, 'BUNIT', ext=1) != 'ELECTRONS'):
            LOG.warn(image + ' is not in ELECTRONS. Skipping...')
            continue

        # Skip processing CTECORR-ed images
        if (fits.getval(image, 'PCTECORR') == 'COMPLETE'):
            LOG.warn(image + ' already has PCTECORR applied. Skipping...')
            continue

        # generate output filename for each input based on specification
        # of the output suffix
        output = image.replace('_flt', '_flt_' + suffix)
        LOG.info('Processing ' + image)
        perform_correction(image, output, maxiter=maxiter, sigrej=sigrej,
                           clobber=clobber)
        LOG.info(output + ' created')

def perform_correction(image, output, maxiter=15, sigrej=2.0, clobber=False):
    """
    Called by `clean` for each input image.

    Parameters
    ----------
    image : str
        Input image name.

    output : str
        Output image name.

    maxiter, sigrej, clobber : see `clean`

    """

    # construct the frame to be cleaned, including the
    # associated data stuctures needed for cleaning
    frame = StripeArray(image)

    # Do the stripe cleaning
    clean_streak(frame, maxiter=maxiter, sigrej=sigrej)

    frame.write_corrected(output, clobber=clobber)

def clean_streak(image, maxiter=15, sigrej=2.0):
    """
    Apply destriping algorithm to input array.

    Parameters
    ----------
    image : `StripeArray` object
        Arrays are modifed in-place.

    maxiter, sigrej : see `clean`

    """

    # create the array to hold the stripe amplitudes
    corr = np.empty(image.science.shape[0])

    # loop over rows to fit the stripe amplitude
    for i in range(image.science.shape[0]):
        # row-by-row iterative sigma-clipped mean; sigma, iters are adjustable
        SMean, SSig, SMedian, SMask = djs_iterstat(
            image.science[i], MaxIter=maxiter, SigRej=sigrej)

        # SExtractor-esque central value statistic; slightly sturdier against
        # skewness of pixel histogram due to faint source flux
        corr[i] = 2.5 * SMedian - 1.5 * SMean

    # preserve the original mean level of the image
    corr -= np.average(corr)

    # apply the correction row-by-row
    for i in range(image.science.shape[0]):
        # stripe is constant along the row, before flatfielding;
        # afterwards it has the shape of the inverse flatfield
        truecorr = corr[i] * image.invflat[i] / np.average(image.invflat[i])

        # correct the SCI extension
        image.science[i] -= truecorr

        # correct the ERR extension
        image.err[i] = np.sqrt(np.abs(image.err[i]**2 - truecorr))

def djs_iterstat(InputArr, MaxIter=10, SigRej=3.0, Max=None, Min=None,
                 RejVal=None):
    """
    Iterative sigma-clipping.

    Parameters
    ----------
    InputArr : `numpy.ndarray`
        Input image array.

    MaxIter, SigRej : see `clean`

    Max, Min : float
        Max and min values for clipping.

    RejVal : float
        Array value to reject, in addition to clipping.

    Returns
    -------
    FMean, FSig, FMedian : float
        Mean, sigma, and median of final result.

    SaveMask : `numpy.ndarray`
        Image mask from the final iteration.

    """
    NGood    = InputArr.size
    ArrShape = InputArr.shape
    if NGood == 0:
        LOG.warn('djs_iterstat - No data points given')
        return 0, 0, 0, 0
    if NGood == 1:
        LOG.warn('djs_iterstat - Only one data point; cannot compute stats')
        return 0, 0, 0, 0
    if np.unique(InputArr).size == 1:
        LOG.warn('djs_iterstat - Only one value in data; cannot compute stats')
        return 0, 0, 0, 0

    # Determine Max and Min
    if Max is None:
        Max = InputArr.max()
    if Min is None:
        Min = InputArr.min()

    Mask = np.zeros(ArrShape, dtype=np.byte) + 1

    # Reject those above Max and those below Min
    Mask[InputArr > Max] = 0
    Mask[InputArr < Min] = 0
    if RejVal is not None:
        Mask[InputArr == RejVal] = 0
    FMean = np.sum(1.0 * InputArr * Mask) / NGood
    FSig  = np.sqrt(np.sum((1.0 * InputArr - FMean)**2 * Mask) / (NGood - 1))

    NLast = -1
    Iter  =  0
    NGood = np.sum(Mask)
    if NGood < 2:
        LOG.warn('djs_iterstat - No good data points; cannot compute stats')
        return -1, -1, -1, -1

    while (Iter < MaxIter) and (NLast != NGood) and (NGood >= 2):
        LoVal = FMean - SigRej * FSig
        HiVal = FMean + SigRej * FSig
        NLast = NGood

        Mask[InputArr < LoVal] = 0
        Mask[InputArr > HiVal] = 0
        NGood = np.sum(Mask)

        if NGood >= 2:
            FMean = np.sum(1.0 * InputArr * Mask) / NGood
            FSig  = np.sqrt(np.sum(
                (1.0 * InputArr - FMean)**2 * Mask) / (NGood - 1))

        SaveMask = Mask.copy()

        Iter += 1

    if np.sum(SaveMask) > 2:
        FMedian = np.median(InputArr[SaveMask == 1])
    else:
        FMedian = FMean

    return FMean, FSig, FMedian, SaveMask

#-------------------------#
# Interfaces used by TEAL #
#-------------------------#

def getHelpAsString(fulldoc=True):
    """
    Returns documentation on the `clean` function. Required by TEAL.

    """
    return clean.__doc__

def run(configobj=None):
    """
    TEAL interface for the `clean` function.

    """
    clean(configobj['input'],
          configobj['suffix'],
          maxiter=configobj['maxiter'],
          sigrej=configobj['sigrej'],
          clobber=configobj['clobber'])

#-----------------------------#
# Interfaces for command line #
#-----------------------------#

def main():
    """Command line driver."""
    import getopt

    usg_str = 'USAGE: acs_destripe [-h][-c] input output [maxiter # [sigrej #]]'

    try:
        optlist, args = getopt.getopt(sys.argv[1:], 'hc')
    except getopt.error as e:
      print(str(e))
      print(usg_str)
      print('\t', __version__)

    # initialize default values
    help = False
    clobber = False
    maxiter = 15
    sigrej = 2.0

    # read options
    for opt, value in optlist:
        if opt == '-h':
            help = True
        elif opt == '-c':
            clobber = True

    if help:
        print(usg_str)
        print('\t', __version__, '(', __vdate__, ')')
        return

    if len(args) < 2:
        LOG.error(usg_str)
        return

    if len(args) > 2:
        # User provided parameters for maxiter, and possibly sigrej
        maxiter = int(args[2])
        if len(args) == 4:
            sigrej = float(args[3])

    clean(args[0], args[1], clobber=clobber, maxiter=maxiter, sigrej=sigrej)

if __name__ == '__main__':
    main()
<commit_msg>Increased acs_destripe version number to reflect r27468<commit_after>#!/usr/bin/env python
"""
Remove horizontal stripes from ACS WFC post-SM4 data.

Examples
--------

In Python without TEAL:

>>> from acstools import acs_destripe
>>> acs_destripe.clean('uncorrected_flt.fits', 'csck', clobber=False,
...                    maxiter=15, sigrej=2.0)

In Python with TEAL:

>>> from acstools import acs_destripe
>>> from stsci.tools import teal
>>> teal.teal('acs_destripe')

In Pyraf::

    --> import acstools
    --> teal acs_destripe

From command line::

    % ./acs_destripe [-h][-c] input output [maxiter # [sigrej #]]

"""
from __future__ import print_function

# STDLIB
import logging
import os
import sys

# THIRD-PARTY
import numpy as np
from astropy.io import fits

# STSCI
from stsci.tools import parseinput,teal

__taskname__ = 'acs_destripe'
__version__ = '0.3.0'
__vdate__ = '18-Nov-2013'
__author__ = 'Norman Grogin, STScI, March 2012.'

MJD_SM4 = 54967

logging.basicConfig()
LOG = logging.getLogger(__taskname__)
LOG.setLevel(logging.INFO)

class StripeArray(object):
    """Class to handle data array to be destriped."""

    def __init__(self, image):
        self.hdulist = fits.open(image)
        self.ampstring = self.hdulist[0].header['CCDAMP']
        self.flatcorr = self.hdulist[0].header['FLATCORR']
        self.configure_arrays()

    def configure_arrays(self):
        """Get the SCI and ERR data."""
        self.science = self.hdulist['sci',1].data
        self.err = self.hdulist['err',1].data
        if (self.ampstring == 'ABCD'):
            self.science = np.concatenate(
                (self.science, self.hdulist['sci',2].data[::-1,:]), axis=1)
            self.err = np.concatenate(
                (self.err, self.hdulist['err',2].data[::-1,:]), axis=1)
        self.ingest_flatfield()

    def ingest_flatfield(self):
        """Process flatfield."""

        flatfile = self.hdulist[0].header['PFLTFILE']

        # if BIAS or DARK, set flatfield to unity
        if flatfile == 'N/A':
            self.invflat = np.ones(self.science.shape)
            return

        hduflat = self.resolve_flatname(flatfile)

        if (self.ampstring == 'ABCD'):
            self.invflat = np.concatenate(
                (1 / hduflat['sci',1].data,
                 1 / hduflat['sci',2].data[::-1,:]), axis=1)
        else:
            # complex algorithm to determine proper subarray of flatfield to use

            # which amp?
            if (self.ampstring == 'A' or self.ampstring == 'B' or
                    self.ampstring == 'AB'):
                self.invflat = 1 / hduflat['sci',2].data
            else:
                self.invflat = 1 / hduflat['sci',1].data

            # now, which section?
            sizaxis1 = self.hdulist[1].header['SIZAXIS1']
            sizaxis2 = self.hdulist[1].header['SIZAXIS2']
            centera1 = self.hdulist[1].header['CENTERA1']
            centera2 = self.hdulist[1].header['CENTERA2']

            flataxis1 = hduflat[1].header['NAXIS1']
            flataxis2 = hduflat[1].header['NAXIS2']

            # configure the offset appropriate to left- or right-side of CCD
            if (self.ampstring[0] == 'A' or self.ampstring[0] == 'C'):
                xdelta = 13
            else:
                xdelta = 35

            if sizaxis1 == flataxis1:
                xlo = 0
                xhi = sizaxis1
            else:
                xlo = centera1 - xdelta - sizaxis1 / 2 - 1
                xhi = centera1 - xdelta + sizaxis1 / 2 - 1

            if sizaxis2 == flataxis2:
                ylo = 0
                yhi = sizaxis2
            else:
                ylo = centera2 - sizaxis2 / 2 - 1
                yhi = centera2 + sizaxis2 / 2 - 1

            self.invflat = self.invflat[ylo:yhi,xlo:xhi]

        # apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science / self.invflat
            self.err = self.err / self.invflat

    def resolve_flatname(self, flatfile):
        """Resolve the filename into an absolute pathname (if necessary)."""
        flatparts = flatfile.partition('$')
        if flatparts[1] == '$':
            flatdir = os.getenv(flatparts[0])
            if flatdir is None:
                errstr = 'Environment variable ' + flatparts[0] + ' undefined!'
                raise ValueError(errstr)
            flatfile = flatdir + flatparts[2]
        if not os.path.exists(flatfile):
            errstr = 'Flat-field file ' + flatfile + ' could not be found!'
            raise IOError, errstr

        return(fits.open(flatfile))

    def write_corrected(self, output, clobber=False):
        """Write out the destriped data."""

        # un-apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science * self.invflat
            self.err = self.err * self.invflat

        # reverse the amp merge
        if (self.ampstring == 'ABCD'):
            tmp_1, tmp_2 = np.split(self.science, 2, axis=1)
            self.hdulist['sci',1].data = tmp_1.copy()
            self.hdulist['sci',2].data = tmp_2[::-1,:].copy()

            tmp_1, tmp_2 = np.split(self.err, 2, axis=1)
            self.hdulist['err',1].data = tmp_1.copy()
            self.hdulist['err',2].data = tmp_2[::-1,:].copy()

        # Write the output
        self.hdulist.writeto(output, clobber=clobber)

def clean(input, suffix, maxiter=15, sigrej=2.0, clobber=False):
    """
    Remove horizontal stripes from ACS WFC post-SM4 data.

    .. note::

        Input data must be an ACS/WFC FLT image, with 2 SCI extensions.
        Does not work on RAW image.

        Uses the flatfield specified by the image header keyword PFLTFILE.
        If keyword value is 'N/A', as is the case with biases and darks,
        then unity flatfield is used.

    Parameters
    ----------
    input : str or list of str
        Input filenames in one of these formats:

            * a Python list of filenames
            * a partial filename with wildcards ('\*flt.fits')
            * filename of an ASN table ('j12345670_asn.fits')
            * an at-file (``@input``)

    suffix : str
        The string to use to add to each input file name to
        indicate an output product. This string will be appended
        to the _flt suffix in each input filename to create the
        new output filename. For example, setting `suffix='csck'`
        will create '\*_flt_csck.fits' images.

    maxiter : int
        This parameter controls the maximum number of iterations
        to perform when computing the statistics used to compute the
        row-by-row corrections.

    sigrej : float
        This parameters sets the sigma level for the rejection applied
        during each iteration of statistics computations for the
        row-by-row corrections.

    clobber : bool
        Specify whether or not to 'clobber' (delete then replace)
        previously generated products with the same names.

    """
    flist, alist = parseinput.parseinput(input)

    for image in flist:
        # Skip processing pre-SM4 images
        if (fits.getval(image, 'EXPSTART') <= MJD_SM4):
            LOG.warn(image + ' is pre-SM4. Skipping...'%image)
            continue

        # Skip processing non-DARKCORR-ed non-BIAS images
        if (fits.getval(image, 'IMAGETYP') != 'BIAS' and
                fits.getval(image,'DARKCORR') != 'COMPLETE'):
            LOG.warn(image + ' is not BIAS and does not have '
                         'DARKCORR=COMPLETE. Skipping...')
            continue

        # Data must be in ELECTRONS
        if (fits.getval(image, 'BUNIT', ext=1) != 'ELECTRONS'):
            LOG.warn(image + ' is not in ELECTRONS. Skipping...')
            continue

        # Skip processing CTECORR-ed images
        if (fits.getval(image, 'PCTECORR') == 'COMPLETE'):
            LOG.warn(image + ' already has PCTECORR applied. Skipping...')
            continue

        # generate output filename for each input based on specification
        # of the output suffix
        output = image.replace('_flt', '_flt_' + suffix)
        LOG.info('Processing ' + image)
        perform_correction(image, output, maxiter=maxiter, sigrej=sigrej,
                           clobber=clobber)
        LOG.info(output + ' created')

def perform_correction(image, output, maxiter=15, sigrej=2.0, clobber=False):
    """
    Called by `clean` for each input image.

    Parameters
    ----------
    image : str
        Input image name.

    output : str
        Output image name.

    maxiter, sigrej, clobber : see `clean`

    """

    # construct the frame to be cleaned, including the
    # associated data stuctures needed for cleaning
    frame = StripeArray(image)

    # Do the stripe cleaning
    clean_streak(frame, maxiter=maxiter, sigrej=sigrej)

    frame.write_corrected(output, clobber=clobber)

def clean_streak(image, maxiter=15, sigrej=2.0):
    """
    Apply destriping algorithm to input array.

    Parameters
    ----------
    image : `StripeArray` object
        Arrays are modifed in-place.

    maxiter, sigrej : see `clean`

    """

    # create the array to hold the stripe amplitudes
    corr = np.empty(image.science.shape[0])

    # loop over rows to fit the stripe amplitude
    for i in range(image.science.shape[0]):
        # row-by-row iterative sigma-clipped mean; sigma, iters are adjustable
        SMean, SSig, SMedian, SMask = djs_iterstat(
            image.science[i], MaxIter=maxiter, SigRej=sigrej)

        # SExtractor-esque central value statistic; slightly sturdier against
        # skewness of pixel histogram due to faint source flux
        corr[i] = 2.5 * SMedian - 1.5 * SMean

    # preserve the original mean level of the image
    corr -= np.average(corr)

    # apply the correction row-by-row
    for i in range(image.science.shape[0]):
        # stripe is constant along the row, before flatfielding;
        # afterwards it has the shape of the inverse flatfield
        truecorr = corr[i] * image.invflat[i] / np.average(image.invflat[i])

        # correct the SCI extension
        image.science[i] -= truecorr

        # correct the ERR extension
        image.err[i] = np.sqrt(np.abs(image.err[i]**2 - truecorr))

def djs_iterstat(InputArr, MaxIter=10, SigRej=3.0, Max=None, Min=None,
                 RejVal=None):
    """
    Iterative sigma-clipping.

    Parameters
    ----------
    InputArr : `numpy.ndarray`
        Input image array.

    MaxIter, SigRej : see `clean`

    Max, Min : float
        Max and min values for clipping.

    RejVal : float
        Array value to reject, in addition to clipping.

    Returns
    -------
    FMean, FSig, FMedian : float
        Mean, sigma, and median of final result.

    SaveMask : `numpy.ndarray`
        Image mask from the final iteration.

    """
    NGood    = InputArr.size
    ArrShape = InputArr.shape
    if NGood == 0:
        LOG.warn('djs_iterstat - No data points given')
        return 0, 0, 0, 0
    if NGood == 1:
        LOG.warn('djs_iterstat - Only one data point; cannot compute stats')
        return 0, 0, 0, 0
    if np.unique(InputArr).size == 1:
        LOG.warn('djs_iterstat - Only one value in data; cannot compute stats')
        return 0, 0, 0, 0

    # Determine Max and Min
    if Max is None:
        Max = InputArr.max()
    if Min is None:
        Min = InputArr.min()

    Mask = np.zeros(ArrShape, dtype=np.byte) + 1

    # Reject those above Max and those below Min
    Mask[InputArr > Max] = 0
    Mask[InputArr < Min] = 0
    if RejVal is not None:
        Mask[InputArr == RejVal] = 0
    FMean = np.sum(1.0 * InputArr * Mask) / NGood
    FSig  = np.sqrt(np.sum((1.0 * InputArr - FMean)**2 * Mask) / (NGood - 1))

    NLast = -1
    Iter  =  0
    NGood = np.sum(Mask)
    if NGood < 2:
        LOG.warn('djs_iterstat - No good data points; cannot compute stats')
        return -1, -1, -1, -1

    while (Iter < MaxIter) and (NLast != NGood) and (NGood >= 2):
        LoVal = FMean - SigRej * FSig
        HiVal = FMean + SigRej * FSig
        NLast = NGood

        Mask[InputArr < LoVal] = 0
        Mask[InputArr > HiVal] = 0
        NGood = np.sum(Mask)

        if NGood >= 2:
            FMean = np.sum(1.0 * InputArr * Mask) / NGood
            FSig  = np.sqrt(np.sum(
                (1.0 * InputArr - FMean)**2 * Mask) / (NGood - 1))

        SaveMask = Mask.copy()

        Iter += 1

    if np.sum(SaveMask) > 2:
        FMedian = np.median(InputArr[SaveMask == 1])
    else:
        FMedian = FMean

    return FMean, FSig, FMedian, SaveMask

#-------------------------#
# Interfaces used by TEAL #
#-------------------------#

def getHelpAsString(fulldoc=True):
    """
    Returns documentation on the `clean` function. Required by TEAL.

    """
    return clean.__doc__

def run(configobj=None):
    """
    TEAL interface for the `clean` function.

    """
    clean(configobj['input'],
          configobj['suffix'],
          maxiter=configobj['maxiter'],
          sigrej=configobj['sigrej'],
          clobber=configobj['clobber'])

#-----------------------------#
# Interfaces for command line #
#-----------------------------#

def main():
    """Command line driver."""
    import getopt

    usg_str = 'USAGE: acs_destripe [-h][-c] input output [maxiter # [sigrej #]]'

    try:
        optlist, args = getopt.getopt(sys.argv[1:], 'hc')
    except getopt.error as e:
      print(str(e))
      print(usg_str)
      print('\t', __version__)

    # initialize default values
    help = False
    clobber = False
    maxiter = 15
    sigrej = 2.0

    # read options
    for opt, value in optlist:
        if opt == '-h':
            help = True
        elif opt == '-c':
            clobber = True

    if help:
        print(usg_str)
        print('\t', __version__, '(', __vdate__, ')')
        return

    if len(args) < 2:
        LOG.error(usg_str)
        return

    if len(args) > 2:
        # User provided parameters for maxiter, and possibly sigrej
        maxiter = int(args[2])
        if len(args) == 4:
            sigrej = float(args[3])

    clean(args[0], args[1], clobber=clobber, maxiter=maxiter, sigrej=sigrej)

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


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

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

import itertools
import operator
import uuid

from functools import partial
from inspect import getmembers
from io import FileIO

from six import iteritems, string_types

from jinja2.exceptions import UndefinedError

from ansible.errors import AnsibleParserError
from ansible.parsing import DataLoader
from ansible.playbook.attribute import Attribute, FieldAttribute
from ansible.template import Templar
from ansible.utils.boolean import boolean
from ansible.utils.debug import debug
from ansible.utils.vars import combine_vars, isidentifier
from ansible.template import template

class Base:

    # connection/transport
    _connection          = FieldAttribute(isa='string')
    _port                = FieldAttribute(isa='int')
    _remote_user         = FieldAttribute(isa='string')

    # variables
    _vars                = FieldAttribute(isa='dict', default=dict())

    # flags and misc. settings
    _environment         = FieldAttribute(isa='list')
    _no_log              = FieldAttribute(isa='bool')

    def __init__(self):

        # initialize the data loader and variable manager, which will be provided
        # later when the object is actually loaded
        self._loader = None
        self._variable_manager = None

        # every object gets a random uuid:
        self._uuid = uuid.uuid4()

        # and initialize the base attributes
        self._initialize_base_attributes()

        try:
            from __main__ import display
            self._display = display
        except ImportError:
            from ansible.utils.display import Display
            self._display = Display()

    # The following three functions are used to programatically define data
    # descriptors (aka properties) for the Attributes of all of the playbook
    # objects (tasks, blocks, plays, etc).
    #
    # The function signature is a little strange because of how we define
    # them.  We use partial to give each method the name of the Attribute that
    # it is for.  Since partial prefills the positional arguments at the
    # beginning of the function we end up with the first positional argument
    # being allocated to the name instead of to the class instance (self) as
    # normal.  To deal with that we make the property name field the first
    # positional argument and self the second arg.
    #
    # Because these methods are defined inside of the class, they get bound to
    # the instance when the object is created.  After we run partial on them
    # and put the result back into the class as a property, they get bound
    # a second time.  This leads to self being  placed in the arguments twice.
    # To work around that, we mark the functions as @staticmethod so that the
    # first binding to the instance doesn't happen.

    @staticmethod
    def _generic_g(prop_name, self):
        method = "_get_attr_%s" % prop_name
        if hasattr(self, method):
            return getattr(self, method)()

        return self._attributes[prop_name]

    @staticmethod
    def _generic_s(prop_name, self, value):
        self._attributes[prop_name] = value

    @staticmethod
    def _generic_d(prop_name, self):
        del self._attributes[prop_name]

    def _get_base_attributes(self):
        '''
        Returns the list of attributes for this class (or any subclass thereof).
        If the attribute name starts with an underscore, it is removed
        '''
        base_attributes = dict()
        for (name, value) in getmembers(self.__class__):
            if isinstance(value, Attribute):
               if name.startswith('_'):
                   name = name[1:]
               base_attributes[name] = value
        return base_attributes

    def _initialize_base_attributes(self):
        # each class knows attributes set upon it, see Task.py for example
        self._attributes = dict()

        for (name, value) in self._get_base_attributes().items():
            getter = partial(self._generic_g, name)
            setter = partial(self._generic_s, name)
            deleter = partial(self._generic_d, name)

            # Place the property into the class so that cls.name is the
            # property functions.
            setattr(Base, name, property(getter, setter, deleter))

            # Place the value into the instance so that the property can
            # process and hold that value/
            setattr(self, name, value.default)

    def preprocess_data(self, ds):
        ''' infrequently used method to do some pre-processing of legacy terms '''

        for base_class in self.__class__.mro():
            method = getattr(self, "_preprocess_data_%s" % base_class.__name__.lower(), None)
            if method:
                return method(ds)
        return ds

    def load_data(self, ds, variable_manager=None, loader=None):
        ''' walk the input datastructure and assign any values '''

        assert ds is not None

        # the variable manager class is used to manage and merge variables
        # down to a single dictionary for reference in templating, etc.
        self._variable_manager = variable_manager

        # the data loader class is used to parse data from strings and files
        if loader is not None:
            self._loader = loader
        else:
            self._loader = DataLoader()

        # call the preprocess_data() function to massage the data into
        # something we can more easily parse, and then call the validation
        # function on it to ensure there are no incorrect key values
        ds = self.preprocess_data(ds)
        self._validate_attributes(ds)

        # Walk all attributes in the class. We sort them based on their priority
        # so that certain fields can be loaded before others, if they are dependent.
        # FIXME: we currently don't do anything with private attributes but
        #        may later decide to filter them out of 'ds' here.
        base_attributes = self._get_base_attributes()
        for name, attr in sorted(base_attributes.items(), key=operator.itemgetter(1)):
            # copy the value over unless a _load_field method is defined
            if name in ds:
                method = getattr(self, '_load_%s' % name, None)
                if method:
                    self._attributes[name] = method(name, ds[name])
                else:
                    self._attributes[name] = ds[name]

        # run early, non-critical validation
        self.validate()

        # cache the datastructure internally
        setattr(self, '_ds', ds)

        # return the constructed object
        return self

    def get_ds(self):
        try:
            return getattr(self, '_ds')
        except AttributeError:
            return None

    def get_loader(self):
        return self._loader

    def get_variable_manager(self):
        return self._variable_manager

    def _validate_attributes(self, ds):
        '''
        Ensures that there are no keys in the datastructure which do
        not map to attributes for this object.
        '''

        valid_attrs = frozenset(name for name in self._get_base_attributes())
        for key in ds:
            if key not in valid_attrs:
                raise AnsibleParserError("'%s' is not a valid attribute for a %s" % (key, self.__class__.__name__), obj=ds)

    def validate(self, all_vars=dict()):
        ''' validation that is done at parse time, not load time '''

        # walk all fields in the object
        for (name, attribute) in iteritems(self._get_base_attributes()):

            # run validator only if present
            method = getattr(self, '_validate_%s' % name, None)
            if method:
                method(attribute, name, getattr(self, name))

    def copy(self):
        '''
        Create a copy of this object and return it.
        '''

        new_me = self.__class__()

        for name in self._get_base_attributes():
            setattr(new_me, name, getattr(self, name))

        new_me._loader           = self._loader
        new_me._variable_manager = self._variable_manager

        # if the ds value was set on the object, copy it to the new copy too
        if hasattr(self, '_ds'):
            new_me._ds = self._ds

        return new_me

    def post_validate(self, templar):
        '''
        we can't tell that everything is of the right type until we have
        all the variables.  Run basic types (from isa) as well as
        any _post_validate_<foo> functions.
        '''

        basedir = None
        if self._loader is not None:
            basedir = self._loader.get_basedir()

        # save the omit value for later checking
        omit_value = templar._available_variables.get('omit')

        for (name, attribute) in iteritems(self._get_base_attributes()):

            if getattr(self, name) is None:
                if not attribute.required:
                    continue
                else:
                    raise AnsibleParserError("the field '%s' is required but was not set" % name)

            try:
                # Run the post-validator if present. These methods are responsible for
                # using the given templar to template the values, if required.
                method = getattr(self, '_post_validate_%s' % name, None)
                if method:
                    value = method(attribute, getattr(self, name), templar)
                else:
                    # if the attribute contains a variable, template it now
                    value = templar.template(getattr(self, name))

                # if this evaluated to the omit value, set the value back to
                # the default specified in the FieldAttribute and move on
                if omit_value is not None and value == omit_value:
                    value = attribute.default
                    continue

                # and make sure the attribute is of the type it should be
                if value is not None:
                    if attribute.isa == 'string':
                        value = unicode(value)
                    elif attribute.isa == 'int':
                        value = int(value)
                    elif attribute.isa == 'float':
                        value = float(value)
                    elif attribute.isa == 'bool':
                        value = boolean(value)
                    elif attribute.isa == 'percent':
                        # special value, which may be an integer or float
                        # with an optional '%' at the end
                        if isinstance(value, string_types) and '%' in value:
                            value = value.replace('%', '')
                        value = float(value)
                    elif attribute.isa == 'list':
                        if value is None:
                            value = []
                        elif not isinstance(value, list):
                            value = [ value ]
                        if attribute.listof is not None:
                            for item in value:
                                if not isinstance(item, attribute.listof):
                                    raise AnsibleParserError("the field '%s' should be a list of %s, but the item '%s' is a %s" % (name, attribute.listof, item, type(item)), obj=self.get_ds())
                                elif attribute.required and attribute.listof == string_types:
                                    if item is None or item.strip() == "":
                                        raise AnsibleParserError("the field '%s' is required, and cannot have empty values" % (name,), obj=self.get_ds())
                    elif attribute.isa == 'set':
                        if value is None:
                            value = set()
                        else:
                            if not isinstance(value, (list, set)):
                                value = [ value ]
                            if not isinstance(value, set):
                                value = set(value)
                    elif attribute.isa == 'dict':
                        if value is None:
                            value = dict()
                        elif not isinstance(value, dict):
                            raise TypeError("%s is not a dictionary" % value)

                # and assign the massaged value back to the attribute field
                setattr(self, name, value)

            except (TypeError, ValueError) as e:
                raise AnsibleParserError("the field '%s' has an invalid value (%s), and could not be converted to an %s. Error was: %s" % (name, value, attribute.isa, e), obj=self.get_ds())
            except UndefinedError as e:
                if templar._fail_on_undefined_errors and name != 'name':
                    raise AnsibleParserError("the field '%s' has an invalid value, which appears to include a variable that is undefined. The error was: %s" % (name,e), obj=self.get_ds())

    def serialize(self):
        '''
        Serializes the object derived from the base object into
        a dictionary of values. This only serializes the field
        attributes for the object, so this may need to be overridden
        for any classes which wish to add additional items not stored
        as field attributes.
        '''

        repr = dict()

        for name in self._get_base_attributes():
            repr[name] = getattr(self, name)

        # serialize the uuid field
        repr['uuid'] = getattr(self, '_uuid')

        return repr

    def deserialize(self, data):
        '''
        Given a dictionary of values, load up the field attributes for
        this object. As with serialize(), if there are any non-field
        attribute data members, this method will need to be overridden
        and extended.
        '''

        assert isinstance(data, dict)

        for (name, attribute) in iteritems(self._get_base_attributes()):
            if name in data:
                setattr(self, name, data[name])
            else:
                setattr(self, name, attribute.default)

        # restore the UUID field
        setattr(self, '_uuid', data.get('uuid'))

    def _load_vars(self, attr, ds):
        '''
        Vars in a play can be specified either as a dictionary directly, or
        as a list of dictionaries. If the later, this method will turn the
        list into a single dictionary.
        '''

        def _validate_variable_keys(ds):
            for key in ds:
                if not isidentifier(key):
                    raise TypeError("%s is not a valid variable name" % key)

        try:
            if isinstance(ds, dict):
                _validate_variable_keys(ds)
                return ds
            elif isinstance(ds, list):
                all_vars = dict()
                for item in ds:
                    if not isinstance(item, dict):
                        raise ValueError
                    _validate_variable_keys(item)
                    all_vars = combine_vars(all_vars, item)
                return all_vars
            elif ds is None:
                return {}
            else:
                raise ValueError
        except ValueError:
            raise AnsibleParserError("Vars in a %s must be specified as a dictionary, or a list of dictionaries" % self.__class__.__name__, obj=ds)
        except TypeError, e:
            raise AnsibleParserError("Invalid variable name in vars specified for %s: %s" % (self.__class__.__name__, e), obj=ds)

    def _extend_value(self, value, new_value):
        '''
        Will extend the value given with new_value (and will turn both
        into lists if they are not so already). The values are run through
        a set to remove duplicate values.
        '''

        if not isinstance(value, list):
            value = [ value ]
        if not isinstance(new_value, list):
            new_value = [ new_value ]

        #return list(set(value + new_value))
        return [i for i,_ in itertools.groupby(value + new_value)]

    def __getstate__(self):
        return self.serialize()

    def __setstate__(self, data):
        self.__init__()
        self.deserialize(data)

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

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

import itertools
import operator
import uuid

from functools import partial
from inspect import getmembers
from io import FileIO

from six import iteritems, string_types

from jinja2.exceptions import UndefinedError

from ansible.errors import AnsibleParserError
from ansible.parsing import DataLoader
from ansible.playbook.attribute import Attribute, FieldAttribute
from ansible.template import Templar
from ansible.utils.boolean import boolean
from ansible.utils.debug import debug
from ansible.utils.vars import combine_vars, isidentifier
from ansible.template import template

class Base:

    # connection/transport
    _connection          = FieldAttribute(isa='string')
    _port                = FieldAttribute(isa='int')
    _remote_user         = FieldAttribute(isa='string')

    # variables
    _vars                = FieldAttribute(isa='dict', default=dict())

    # flags and misc. settings
    _environment         = FieldAttribute(isa='list')
    _no_log              = FieldAttribute(isa='bool')

    def __init__(self):

        # initialize the data loader and variable manager, which will be provided
        # later when the object is actually loaded
        self._loader = None
        self._variable_manager = None

        # every object gets a random uuid:
        self._uuid = uuid.uuid4()

        # and initialize the base attributes
        self._initialize_base_attributes()

        try:
            from __main__ import display
            self._display = display
        except ImportError:
            from ansible.utils.display import Display
            self._display = Display()

    # The following three functions are used to programatically define data
    # descriptors (aka properties) for the Attributes of all of the playbook
    # objects (tasks, blocks, plays, etc).
    #
    # The function signature is a little strange because of how we define
    # them.  We use partial to give each method the name of the Attribute that
    # it is for.  Since partial prefills the positional arguments at the
    # beginning of the function we end up with the first positional argument
    # being allocated to the name instead of to the class instance (self) as
    # normal.  To deal with that we make the property name field the first
    # positional argument and self the second arg.
    #
    # Because these methods are defined inside of the class, they get bound to
    # the instance when the object is created.  After we run partial on them
    # and put the result back into the class as a property, they get bound
    # a second time.  This leads to self being  placed in the arguments twice.
    # To work around that, we mark the functions as @staticmethod so that the
    # first binding to the instance doesn't happen.

    @staticmethod
    def _generic_g(prop_name, self):
        method = "_get_attr_%s" % prop_name
        if hasattr(self, method):
            return getattr(self, method)()

        return self._attributes[prop_name]

    @staticmethod
    def _generic_s(prop_name, self, value):
        self._attributes[prop_name] = value

    @staticmethod
    def _generic_d(prop_name, self):
        del self._attributes[prop_name]

    def _get_base_attributes(self):
        '''
        Returns the list of attributes for this class (or any subclass thereof).
        If the attribute name starts with an underscore, it is removed
        '''
        base_attributes = dict()
        for (name, value) in getmembers(self.__class__):
            if isinstance(value, Attribute):
               if name.startswith('_'):
                   name = name[1:]
               base_attributes[name] = value
        return base_attributes

    def _initialize_base_attributes(self):
        # each class knows attributes set upon it, see Task.py for example
        self._attributes = dict()

        for (name, value) in self._get_base_attributes().items():
            getter = partial(self._generic_g, name)
            setter = partial(self._generic_s, name)
            deleter = partial(self._generic_d, name)

            # Place the property into the class so that cls.name is the
            # property functions.
            setattr(Base, name, property(getter, setter, deleter))

            # Place the value into the instance so that the property can
            # process and hold that value/
            setattr(self, name, value.default)

    def preprocess_data(self, ds):
        ''' infrequently used method to do some pre-processing of legacy terms '''

        for base_class in self.__class__.mro():
            method = getattr(self, "_preprocess_data_%s" % base_class.__name__.lower(), None)
            if method:
                return method(ds)
        return ds

    def load_data(self, ds, variable_manager=None, loader=None):
        ''' walk the input datastructure and assign any values '''

        assert ds is not None

        # the variable manager class is used to manage and merge variables
        # down to a single dictionary for reference in templating, etc.
        self._variable_manager = variable_manager

        # the data loader class is used to parse data from strings and files
        if loader is not None:
            self._loader = loader
        else:
            self._loader = DataLoader()

        # call the preprocess_data() function to massage the data into
        # something we can more easily parse, and then call the validation
        # function on it to ensure there are no incorrect key values
        ds = self.preprocess_data(ds)
        self._validate_attributes(ds)

        # Walk all attributes in the class. We sort them based on their priority
        # so that certain fields can be loaded before others, if they are dependent.
        # FIXME: we currently don't do anything with private attributes but
        #        may later decide to filter them out of 'ds' here.
        base_attributes = self._get_base_attributes()
        for name, attr in sorted(base_attributes.items(), key=operator.itemgetter(1)):
            # copy the value over unless a _load_field method is defined
            if name in ds:
                method = getattr(self, '_load_%s' % name, None)
                if method:
                    self._attributes[name] = method(name, ds[name])
                else:
                    self._attributes[name] = ds[name]

        # run early, non-critical validation
        self.validate()

        # cache the datastructure internally
        setattr(self, '_ds', ds)

        # return the constructed object
        return self

    def get_ds(self):
        try:
            return getattr(self, '_ds')
        except AttributeError:
            return None

    def get_loader(self):
        return self._loader

    def get_variable_manager(self):
        return self._variable_manager

    def _validate_attributes(self, ds):
        '''
        Ensures that there are no keys in the datastructure which do
        not map to attributes for this object.
        '''

        valid_attrs = frozenset(name for name in self._get_base_attributes())
        for key in ds:
            if key not in valid_attrs:
                raise AnsibleParserError("'%s' is not a valid attribute for a %s" % (key, self.__class__.__name__), obj=ds)

    def validate(self, all_vars=dict()):
        ''' validation that is done at parse time, not load time '''

        # walk all fields in the object
        for (name, attribute) in iteritems(self._get_base_attributes()):

            # run validator only if present
            method = getattr(self, '_validate_%s' % name, None)
            if method:
                method(attribute, name, getattr(self, name))

    def copy(self):
        '''
        Create a copy of this object and return it.
        '''

        new_me = self.__class__()

        for name in self._get_base_attributes():
            setattr(new_me, name, getattr(self, name))

        new_me._loader           = self._loader
        new_me._variable_manager = self._variable_manager

        # if the ds value was set on the object, copy it to the new copy too
        if hasattr(self, '_ds'):
            new_me._ds = self._ds

        return new_me

    def post_validate(self, templar):
        '''
        we can't tell that everything is of the right type until we have
        all the variables.  Run basic types (from isa) as well as
        any _post_validate_<foo> functions.
        '''

        basedir = None
        if self._loader is not None:
            basedir = self._loader.get_basedir()

        # save the omit value for later checking
        omit_value = templar._available_variables.get('omit')

        for (name, attribute) in iteritems(self._get_base_attributes()):

            if getattr(self, name) is None:
                if not attribute.required:
                    continue
                else:
                    raise AnsibleParserError("the field '%s' is required but was not set" % name)

            try:
                # Run the post-validator if present. These methods are responsible for
                # using the given templar to template the values, if required.
                method = getattr(self, '_post_validate_%s' % name, None)
                if method:
                    value = method(attribute, getattr(self, name), templar)
                else:
                    # if the attribute contains a variable, template it now
                    value = templar.template(getattr(self, name))

                # if this evaluated to the omit value, set the value back to
                # the default specified in the FieldAttribute and move on
                if omit_value is not None and value == omit_value:
                    value = attribute.default
                    continue

                # and make sure the attribute is of the type it should be
                if value is not None:
                    if attribute.isa == 'string':
                        value = unicode(value)
                    elif attribute.isa == 'int':
                        value = int(value)
                    elif attribute.isa == 'float':
                        value = float(value)
                    elif attribute.isa == 'bool':
                        value = boolean(value)
                    elif attribute.isa == 'percent':
                        # special value, which may be an integer or float
                        # with an optional '%' at the end
                        if isinstance(value, string_types) and '%' in value:
                            value = value.replace('%', '')
                        value = float(value)
                    elif attribute.isa == 'list':
                        if value is None:
                            value = []
                        elif not isinstance(value, list):
                            value = [ value ]
                        if attribute.listof is not None:
                            for item in value:
                                if not isinstance(item, attribute.listof):
                                    raise AnsibleParserError("the field '%s' should be a list of %s, but the item '%s' is a %s" % (name, attribute.listof, item, type(item)), obj=self.get_ds())
                                elif attribute.required and attribute.listof == string_types:
                                    if item is None or item.strip() == "":
                                        raise AnsibleParserError("the field '%s' is required, and cannot have empty values" % (name,), obj=self.get_ds())
                    elif attribute.isa == 'set':
                        if value is None:
                            value = set()
                        else:
                            if not isinstance(value, (list, set)):
                                value = [ value ]
                            if not isinstance(value, set):
                                value = set(value)
                    elif attribute.isa == 'dict':
                        if value is None:
                            value = dict()
                        elif not isinstance(value, dict):
                            raise TypeError("%s is not a dictionary" % value)

                # and assign the massaged value back to the attribute field
                setattr(self, name, value)

            except (TypeError, ValueError) as e:
                raise AnsibleParserError("the field '%s' has an invalid value (%s), and could not be converted to an %s. Error was: %s" % (name, value, attribute.isa, e), obj=self.get_ds())
            except UndefinedError as e:
                if templar._fail_on_undefined_errors and name != 'name':
                    raise AnsibleParserError("the field '%s' has an invalid value, which appears to include a variable that is undefined. The error was: %s" % (name,e), obj=self.get_ds())

    def serialize(self):
        '''
        Serializes the object derived from the base object into
        a dictionary of values. This only serializes the field
        attributes for the object, so this may need to be overridden
        for any classes which wish to add additional items not stored
        as field attributes.
        '''

        repr = dict()

        for name in self._get_base_attributes():
            repr[name] = getattr(self, name)

        # serialize the uuid field
        repr['uuid'] = getattr(self, '_uuid')

        return repr

    def deserialize(self, data):
        '''
        Given a dictionary of values, load up the field attributes for
        this object. As with serialize(), if there are any non-field
        attribute data members, this method will need to be overridden
        and extended.
        '''

        assert isinstance(data, dict)

        for (name, attribute) in iteritems(self._get_base_attributes()):
            if name in data:
                setattr(self, name, data[name])
            else:
                setattr(self, name, attribute.default)

        # restore the UUID field
        setattr(self, '_uuid', data.get('uuid'))

    def _load_vars(self, attr, ds):
        '''
        Vars in a play can be specified either as a dictionary directly, or
        as a list of dictionaries. If the later, this method will turn the
        list into a single dictionary.
        '''

        def _validate_variable_keys(ds):
            for key in ds:
                if not isidentifier(key):
                    raise TypeError("%s is not a valid variable name" % key)

        try:
            if isinstance(ds, dict):
                _validate_variable_keys(ds)
                return ds
            elif isinstance(ds, list):
                all_vars = dict()
                for item in ds:
                    if not isinstance(item, dict):
                        raise ValueError
                    _validate_variable_keys(item)
                    all_vars = combine_vars(all_vars, item)
                return all_vars
            elif ds is None:
                return {}
            else:
                raise ValueError
        except ValueError:
            raise AnsibleParserError("Vars in a %s must be specified as a dictionary, or a list of dictionaries" % self.__class__.__name__, obj=ds)
        except TypeError as e:
            raise AnsibleParserError("Invalid variable name in vars specified for %s: %s" % (self.__class__.__name__, e), obj=ds)

    def _extend_value(self, value, new_value):
        '''
        Will extend the value given with new_value (and will turn both
        into lists if they are not so already). The values are run through
        a set to remove duplicate values.
        '''

        if not isinstance(value, list):
            value = [ value ]
        if not isinstance(new_value, list):
            new_value = [ new_value ]

        #return list(set(value + new_value))
        return [i for i,_ in itertools.groupby(value + new_value)]

    def __getstate__(self):
        return self.serialize()

    def __setstate__(self, data):
        self.__init__()
        self.deserialize(data)

<|endoftext|>


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

import os
import os.path
import sys
import glob
import imp
from ansible import constants as C
from ansible import errors

MODULE_CACHE = {}
PATH_CACHE = {}
PLUGIN_PATH_CACHE = {}
_basedirs = []

def push_basedir(basedir):
    if basedir not in _basedirs:
        _basedirs.insert(0, basedir)

class PluginLoader(object):

    '''
    PluginLoader loads plugins from the configured plugin directories.

    It searches for plugins by iterating through the combined list of
    play basedirs, configured paths, and the python path.
    The first match is used.
    '''

    def __init__(self, class_name, package, config, subdir, aliases={}):

        self.class_name         = class_name
        self.package            = package
        self.config             = config
        self.subdir             = subdir
        self.aliases            = aliases

        if not class_name in MODULE_CACHE:
            MODULE_CACHE[class_name] = {}
        if not class_name in PATH_CACHE:
            PATH_CACHE[class_name] = None
        if not class_name in PLUGIN_PATH_CACHE:
            PLUGIN_PATH_CACHE[class_name] = {}

        self._module_cache      = MODULE_CACHE[class_name]
        self._paths             = PATH_CACHE[class_name]
        self._plugin_path_cache = PLUGIN_PATH_CACHE[class_name]

        self._extra_dirs = []

    def print_paths(self):
        ''' Returns a string suitable for printing of the search path '''

        # Uses a list to get the order right
        ret = []
        for i in self._get_paths():
            if i not in ret:
                ret.append(i)
        return os.pathsep.join(ret)

    def _get_package_paths(self):
        ''' Gets the path of a Python package '''

        paths = []
        if not self.package:
            return []
        if not hasattr(self, 'package_path'):
            m = __import__(self.package)
            parts = self.package.split('.')[1:]
            self.package_path = os.path.join(os.path.dirname(m.__file__), *parts)
            paths.append(self.package_path)
            return paths
        else:
            return [ self.package_path ]

    def _get_paths(self):
        ''' Return a list of paths to search for plugins in '''

        if self._paths is not None:
            return self._paths

        ret = []
        ret += self._extra_dirs
        for basedir in _basedirs:
            fullpath = os.path.join(basedir, self.subdir)
            if os.path.isdir(fullpath):
                files = glob.glob("%s/*" % fullpath)
                for file in files:
                    if os.path.isdir(file) and file not in ret:
                        ret.append(file)
                if fullpath not in ret:
                    ret.append(fullpath)

        # look in any configured plugin paths, allow one level deep for subcategories
        configured_paths = self.config.split(os.pathsep)
        for path in configured_paths:
            path = os.path.expanduser(path)
            contents = glob.glob("%s/*" % path)
            for c in contents:
                if os.path.isdir(c):
                    ret.append(c)
            ret.append(path)

        # look for any plugins installed in the package subtree
        ret.extend(self._get_package_paths())

        self._paths = ret

        return ret

    def add_directory(self, directory, with_subdir=False):
        ''' Adds an additional directory to the search path '''

        self._paths = None

        if directory is not None:
            if with_subdir:
                directory = os.path.join(directory, self.subdir)
            if directory not in self._extra_dirs:
                self._extra_dirs.append(directory)

    def find_plugin(self, name):
        ''' Find a plugin named name '''

        if 'name' in self._plugin_path_cache:
            return self._plugin_path_cache[name]

        suffix = ".py"
        if not self.class_name:
            suffix = ""

        for i in self._get_paths():
            path = os.path.join(i, "%s%s" % (name, suffix))
            if os.path.exists(path):
                self._plugin_path_cache[name] = path
                return path

        return None

    def has_plugin(self, name):
        ''' Checks if a plugin named name exists '''

        return self.find_plugin(name) is not None

    __contains__ = has_plugin

    def get(self, name, *args, **kwargs):
        ''' instantiates a plugin of the given name using arguments '''

        if name in self.aliases:
            name = self.aliases[name]
        path = self.find_plugin(name)
        if path is None:
            return None
        if path not in self._module_cache:
            self._module_cache[path] = imp.load_source('.'.join([self.package, name]), path)
        return getattr(self._module_cache[path], self.class_name)(*args, **kwargs)

    def all(self, *args, **kwargs):
        ''' instantiates all plugins with the same arguments '''

        for i in self._get_paths():
            for path in glob.glob(os.path.join(i, "*.py")):
                name, ext = os.path.splitext(os.path.basename(path))
                if name.startswith("_"):
                    continue
                if path not in self._module_cache:
                    self._module_cache[path] = imp.load_source('.'.join([self.package, name]), path)
                yield getattr(self._module_cache[path], self.class_name)(*args, **kwargs)

action_loader = PluginLoader(
    'ActionModule',
    'ansible.runner.action_plugins',
    C.DEFAULT_ACTION_PLUGIN_PATH,
    'action_plugins'
)

callback_loader = PluginLoader(
    'CallbackModule',
    'ansible.callback_plugins',
    C.DEFAULT_CALLBACK_PLUGIN_PATH,
    'callback_plugins'
)

connection_loader = PluginLoader(
    'Connection',
    'ansible.runner.connection_plugins',
    C.DEFAULT_CONNECTION_PLUGIN_PATH,
    'connection_plugins',
    aliases={'paramiko': 'paramiko_ssh'}
)

module_finder = PluginLoader(
    '',
    '',
    C.DEFAULT_MODULE_PATH,
    'library'
)

lookup_loader = PluginLoader(
    'LookupModule',
    'ansible.runner.lookup_plugins',
    C.DEFAULT_LOOKUP_PLUGIN_PATH,
    'lookup_plugins'
)

vars_loader = PluginLoader(
    'VarsModule',
    'ansible.inventory.vars_plugins',
    C.DEFAULT_VARS_PLUGIN_PATH,
    'vars_plugins'
)

filter_loader = PluginLoader(
    'FilterModule',
    'ansible.runner.filter_plugins',
    C.DEFAULT_FILTER_PLUGIN_PATH,
    'filter_plugins'
)

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

import os
import os.path
import sys
import glob
import imp
from ansible import constants as C
from ansible import errors

MODULE_CACHE = {}
PATH_CACHE = {}
PLUGIN_PATH_CACHE = {}
_basedirs = []

def push_basedir(basedir):
    # avoid pushing the same absolute dir more than once
    basedir = os.path.abspath(basedir)
    if basedir not in _basedirs:
        _basedirs.insert(0, basedir)

class PluginLoader(object):

    '''
    PluginLoader loads plugins from the configured plugin directories.

    It searches for plugins by iterating through the combined list of
    play basedirs, configured paths, and the python path.
    The first match is used.
    '''

    def __init__(self, class_name, package, config, subdir, aliases={}):

        self.class_name         = class_name
        self.package            = package
        self.config             = config
        self.subdir             = subdir
        self.aliases            = aliases

        if not class_name in MODULE_CACHE:
            MODULE_CACHE[class_name] = {}
        if not class_name in PATH_CACHE:
            PATH_CACHE[class_name] = None
        if not class_name in PLUGIN_PATH_CACHE:
            PLUGIN_PATH_CACHE[class_name] = {}

        self._module_cache      = MODULE_CACHE[class_name]
        self._paths             = PATH_CACHE[class_name]
        self._plugin_path_cache = PLUGIN_PATH_CACHE[class_name]

        self._extra_dirs = []

    def print_paths(self):
        ''' Returns a string suitable for printing of the search path '''

        # Uses a list to get the order right
        ret = []
        for i in self._get_paths():
            if i not in ret:
                ret.append(i)
        return os.pathsep.join(ret)

    def _get_package_paths(self):
        ''' Gets the path of a Python package '''

        paths = []
        if not self.package:
            return []
        if not hasattr(self, 'package_path'):
            m = __import__(self.package)
            parts = self.package.split('.')[1:]
            self.package_path = os.path.join(os.path.dirname(m.__file__), *parts)
            paths.append(self.package_path)
            return paths
        else:
            return [ self.package_path ]

    def _get_paths(self):
        ''' Return a list of paths to search for plugins in '''

        if self._paths is not None:
            return self._paths

        ret = []
        ret += self._extra_dirs
        for basedir in _basedirs:
            fullpath = os.path.abspath(os.path.join(basedir, self.subdir))
            if os.path.isdir(fullpath):
                files = glob.glob("%s/*" % fullpath)
                for file in files:
                    if os.path.isdir(file) and file not in ret:
                        ret.append(file)
                if fullpath not in ret:
                    ret.append(fullpath)

        # look in any configured plugin paths, allow one level deep for subcategories
        configured_paths = self.config.split(os.pathsep)
        for path in configured_paths:
            path = os.path.abspath(os.path.expanduser(path))
            contents = glob.glob("%s/*" % path)
            for c in contents:
                if os.path.isdir(c) and c not in ret:
                    ret.append(c)
            if path not in ret:
                ret.append(path)

        # look for any plugins installed in the package subtree
        ret.extend(self._get_package_paths())

        self._paths = ret

        return ret

    def add_directory(self, directory, with_subdir=False):
        ''' Adds an additional directory to the search path '''

        self._paths = None
        directory = os.path.abspath(directory)

        if directory is not None:
            if with_subdir:
                directory = os.path.join(directory, self.subdir)
            if directory not in self._extra_dirs:
                self._extra_dirs.append(directory)

    def find_plugin(self, name):
        ''' Find a plugin named name '''

        if 'name' in self._plugin_path_cache:
            return self._plugin_path_cache[name]

        suffix = ".py"
        if not self.class_name:
            suffix = ""

        for i in self._get_paths():
            path = os.path.join(i, "%s%s" % (name, suffix))
            if os.path.exists(path):
                self._plugin_path_cache[name] = path
                return path

        return None

    def has_plugin(self, name):
        ''' Checks if a plugin named name exists '''

        return self.find_plugin(name) is not None

    __contains__ = has_plugin

    def get(self, name, *args, **kwargs):
        ''' instantiates a plugin of the given name using arguments '''

        if name in self.aliases:
            name = self.aliases[name]
        path = self.find_plugin(name)
        if path is None:
            return None
        if path not in self._module_cache:
            self._module_cache[path] = imp.load_source('.'.join([self.package, name]), path)
        return getattr(self._module_cache[path], self.class_name)(*args, **kwargs)

    def all(self, *args, **kwargs):
        ''' instantiates all plugins with the same arguments '''

        for i in self._get_paths():
            for path in glob.glob(os.path.join(i, "*.py")):
                name, ext = os.path.splitext(os.path.basename(path))
                if name.startswith("_"):
                    continue
                if path not in self._module_cache:
                    self._module_cache[path] = imp.load_source('.'.join([self.package, name]), path)
                yield getattr(self._module_cache[path], self.class_name)(*args, **kwargs)

action_loader = PluginLoader(
    'ActionModule',
    'ansible.runner.action_plugins',
    C.DEFAULT_ACTION_PLUGIN_PATH,
    'action_plugins'
)

callback_loader = PluginLoader(
    'CallbackModule',
    'ansible.callback_plugins',
    C.DEFAULT_CALLBACK_PLUGIN_PATH,
    'callback_plugins'
)

connection_loader = PluginLoader(
    'Connection',
    'ansible.runner.connection_plugins',
    C.DEFAULT_CONNECTION_PLUGIN_PATH,
    'connection_plugins',
    aliases={'paramiko': 'paramiko_ssh'}
)

module_finder = PluginLoader(
    '',
    '',
    C.DEFAULT_MODULE_PATH,
    'library'
)

lookup_loader = PluginLoader(
    'LookupModule',
    'ansible.runner.lookup_plugins',
    C.DEFAULT_LOOKUP_PLUGIN_PATH,
    'lookup_plugins'
)

vars_loader = PluginLoader(
    'VarsModule',
    'ansible.inventory.vars_plugins',
    C.DEFAULT_VARS_PLUGIN_PATH,
    'vars_plugins'
)

filter_loader = PluginLoader(
    'FilterModule',
    'ansible.runner.filter_plugins',
    C.DEFAULT_FILTER_PLUGIN_PATH,
    'filter_plugins'
)

<|endoftext|>