<reponame>exoplanetvetting/DAVE<filename>fileio/AbstractMast.py

"""
Abstract classes for interfacing with the Mast Kepler Data archives.
 
Used by mastio.py
"""

import urllib
import math
import os

import pyfits


class MastArchive():
    """Base class for queries of the MAST Kepler and K2 archives.
    Handles submitting the URL, receiving and caching the results

    The child class determines the url to submit
    """
    def __init__(self, localDir, remoteServer):
        self.localDir = localDir
        self.remoteServer = remoteServer

        if not os.path.exists(localDir):
            try:
                os.mkdir(localDir)
            except IOError, e:
                msg = "Local path %s does not exist and can't be created: %s" \
                    %(localDir, e)
                raise IOError(e)


    def getData(self, localUrl, remoteUrl, compressedOnServer, *args, **kwargs):
        forceDownload = kwargs.pop("force", False)

        if not forceDownload:
            if os.path.exists(localUrl):
                isEmpty = os.stat(localUrl)[6]==0
                if isEmpty:
                    raise IOError("File %s previously returned 404" %(remoteUrl))

        if not os.path.exists(localUrl):
            if compressedOnServer:
                self.download(remoteUrl, localUrl + '.gz')
                os.system("gunzip %s" %(localUrl + '.gz'))
            else:
                self.download(remoteUrl, localUrl)

        return self.parse(localUrl, *args, **kwargs)

    
    def download(self, remoteUrl, localUrl, clobber=True, verbose=False):
        """Download data from MAST archive for a given filename and a
        given quarter.

        Inputs:
        clobber:    If true, allow locally existing files to be overwritten
        verbose:    If true, print the file meta-information
        """

        if not clobber:
            if os.path.exists(localUrl):
                raise IOError("File %s already exists. Refusing to overwrite" %(localUrl))

        u = urllib.urlopen(remoteUrl)
        if u.getcode() == 404:
            #Write a zero length file to indicate no file is present
            if localUrl[-3:] == ".gz":
                localUrl = localUrl[:-3]

            fp = open(localUrl, 'w')
            fp.write("")
            fp.close()
            raise ValueError("File not found at MAST: %s" %(remoteUrl))
            return

        try:
            (f,h) = urllib.urlretrieve(remoteUrl, localUrl)
        except IOError, e:
            raise(e)


        if verbose:
            print h
    
    def parse(self, localUrl, *args, **kwargs):
        """Load data using favourite IO class
        
        Assumes that localUrl exists on disk. Default file format
        for MAST is FITS, but sub-classes can override this method
        for other file types as necessary
        """
        if 'ext' in kwargs and kwargs['ext'] == 0:
            try:
                retVal = pyfits.getheader(localUrl, *args, **kwargs)
            except IOError, e:
                raise e
        else:
            try:
                retVal = pyfits.getdata(localUrl, *args, **kwargs)
            except IOError, e:
                raise e

        return retVal




class KeplerAbstractClass(MastArchive):
    """Base class for classes interacting with Kepler/K2 data.
    
    This class contains methods for querying for Lightcurve and 
    TPF files in either short or long cadence.
    
    Example daughter classes are KeplerArchive and K2Archive.
    Daughter classes must, at a minimum, reimplement getFilename()
    and makeRemoteUrl(). You may want to write your own init, too.
    
    """
    
    def __init__(self, localDir, remoteServer, remoteFluxPath, remoteTpfPath):
        MastArchive.__init__(self, localDir, remoteServer)
        self.remoteFluxPath = remoteFluxPath
        self.remoteTpfPath = remoteTpfPath

        
    def getLocalDir(self, quarter, kepid, quarterPrefix="Q"):
        quarterDir = "%s%s" %(quarterPrefix, quarter)

        localDir = os.path.join(self.localDir, quarterDir)
        if not os.path.exists(localDir):
            try:
                os.mkdir(localDir)
            except OSError, e:
                #Re-raising an exception makes the error easier to read
                raise e

        thou = "%04i" %( math.floor( float(kepid)/1000.))
        localDir = os.path.join(localDir, thou)
        if not os.path.exists(localDir):
            try:
                os.makedirs(localDir)
            except OSError. e:
                #Re-raising an exception makes the error easier to read
                raise e

        return localDir


    def getLongCadence(self, kepid, quarter, *args, **kwargs):
        """Get a long cadence lightcurve file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, True, None, *args, **kwargs)


    def getShortCadence(self, kepid, quarter, month, *args, **kwargs):
        """Get a short cadence lightcurve file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, True, month, *args, **kwargs)


    def getLongTpf(self, kepid, quarter, *args, **kwargs):
        """Get a long cadence TPF file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, False, None, *args, **kwargs)


    def getShortTpf(self, kepid, quarter, month, *args, **kwargs):
        """Get a short cadence TPF file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, False, month, *args, **kwargs)


    def getFile(self, kepid, quarter, isFluxFile, month=None, *args, **kwargs):
        """Get a file of a given type for a given kepid/quarter[/month]

        Inputs:
        ---------
        kepid   (int)
            Kepid of target of interest
        quarter (int)
            Quarter of interest
        isFluxFile (bool)
            If True, a lightcurve file is returned, otherwise a target
            pixel file is returned

        Optional Inputs:
        ----------------
        month   (int)
            If None, a long cadence file is returned, if a value
            of 1,2 or 3 is supplied, a short cadence file is returned.

        All other optional inputs are passed to ``pyfits.getdata()``

        Returns:
        -------------
        A FITs object, and possibly more, depending on the optional
        arguments to ``pyfits.getdata``
        """

        filename = self.getFilename(kepid, quarter, isFluxFile, month)
        localUrl = os.path.join( self.getLocalDir(quarter, kepid), filename)

        #Make remote url
        if isFluxFile:
            remotePath = self.remoteFluxPath
            compressedOnServer = False
        else:
            remotePath = self.remoteTpfPath
            compressedOnServer = True

        remoteUrl = self.makeRemoteUrl(remotePath, kepid, quarter, filename, compressedOnServer)
        return self.getData(localUrl, remoteUrl, compressedOnServer, *args, **kwargs)


    def makeRemoteUrl(self, remotePath, kepid, quarter, filename, compressedOnServer):
        """Figure out the URL where the desired file is stored.
        
        Returns a string, typically similar to
        http://archive.stsci.edu/$remotePath/.../$filename
        """
        raise NotImplementedError("Not Implemented in Base Class")

    def getFilename(self, kepid, quarter, isFluxFile, isShortCadence):
        """Figure out the name of the file you want to download
        
        (But not its path)
        """
        raise NotImplementedError("Not Implemented in Base Class")


"""
Abstract classes for interfacing with the Mast Kepler Data archives.
 
Used by mastio.py
"""

import urllib
import math
import os

import pyfits


class MastArchive():
    """Base class for queries of the MAST Kepler and K2 archives.
    Handles submitting the URL, receiving and caching the results

    The child class determines the url to submit
    """
    def __init__(self, localDir, remoteServer):
        self.localDir = localDir
        self.remoteServer = remoteServer

        if not os.path.exists(localDir):
            try:
                os.mkdir(localDir)
            except IOError, e:
                msg = "Local path %s does not exist and can't be created: %s" \
                    %(localDir, e)
                raise IOError(e)


    def getData(self, localUrl, remoteUrl, compressedOnServer, *args, **kwargs):
        forceDownload = kwargs.pop("force", False)

        if not forceDownload:
            if os.path.exists(localUrl):
                isEmpty = os.stat(localUrl)[6]==0
                if isEmpty:
                    raise IOError("File %s previously returned 404" %(remoteUrl))

        if not os.path.exists(localUrl):
            if compressedOnServer:
                self.download(remoteUrl, localUrl + '.gz')
                os.system("gunzip %s" %(localUrl + '.gz'))
            else:
                self.download(remoteUrl, localUrl)

        return self.parse(localUrl, *args, **kwargs)

    
    def download(self, remoteUrl, localUrl, clobber=True, verbose=False):
        """Download data from MAST archive for a given filename and a
        given quarter.

        Inputs:
        clobber:    If true, allow locally existing files to be overwritten
        verbose:    If true, print the file meta-information
        """

        if not clobber:
            if os.path.exists(localUrl):
                raise IOError("File %s already exists. Refusing to overwrite" %(localUrl))

        u = urllib.urlopen(remoteUrl)
        if u.getcode() == 404:
            #Write a zero length file to indicate no file is present
            if localUrl[-3:] == ".gz":
                localUrl = localUrl[:-3]

            fp = open(localUrl, 'w')
            fp.write("")
            fp.close()
            raise ValueError("File not found at MAST: %s" %(remoteUrl))
            return

        try:
            (f,h) = urllib.urlretrieve(remoteUrl, localUrl)
        except IOError, e:
            raise(e)


        if verbose:
            print h
    
    def parse(self, localUrl, *args, **kwargs):
        """Load data using favourite IO class
        
        Assumes that localUrl exists on disk. Default file format
        for MAST is FITS, but sub-classes can override this method
        for other file types as necessary
        """
        if 'ext' in kwargs and kwargs['ext'] == 0:
            try:
                retVal = pyfits.getheader(localUrl, *args, **kwargs)
            except IOError, e:
                raise e
        else:
            try:
                retVal = pyfits.getdata(localUrl, *args, **kwargs)
            except IOError, e:
                raise e

        return retVal




class KeplerAbstractClass(MastArchive):
    """Base class for classes interacting with Kepler/K2 data.
    
    This class contains methods for querying for Lightcurve and 
    TPF files in either short or long cadence.
    
    Example daughter classes are KeplerArchive and K2Archive.
    Daughter classes must, at a minimum, reimplement getFilename()
    and makeRemoteUrl(). You may want to write your own init, too.
    
    """
    
    def __init__(self, localDir, remoteServer, remoteFluxPath, remoteTpfPath):
        MastArchive.__init__(self, localDir, remoteServer)
        self.remoteFluxPath = remoteFluxPath
        self.remoteTpfPath = remoteTpfPath

        
    def getLocalDir(self, quarter, kepid, quarterPrefix="Q"):
        quarterDir = "%s%s" %(quarterPrefix, quarter)

        localDir = os.path.join(self.localDir, quarterDir)
        if not os.path.exists(localDir):
            try:
                os.mkdir(localDir)
            except OSError, e:
                #Re-raising an exception makes the error easier to read
                raise e

        thou = "%04i" %( math.floor( float(kepid)/1000.))
        localDir = os.path.join(localDir, thou)
        if not os.path.exists(localDir):
            try:
                os.makedirs(localDir)
            except OSError. e:
                #Re-raising an exception makes the error easier to read
                raise e

        return localDir


    def getLongCadence(self, kepid, quarter, *args, **kwargs):
        """Get a long cadence lightcurve file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, True, None, *args, **kwargs)


    def getShortCadence(self, kepid, quarter, month, *args, **kwargs):
        """Get a short cadence lightcurve file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, True, month, *args, **kwargs)


    def getLongTpf(self, kepid, quarter, *args, **kwargs):
        """Get a long cadence TPF file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, False, None, *args, **kwargs)


    def getShortTpf(self, kepid, quarter, month, *args, **kwargs):
        """Get a short cadence TPF file.

        A thin wrapper around getFile()
        """
        return self.getFile(kepid, quarter, False, month, *args, **kwargs)


    def getFile(self, kepid, quarter, isFluxFile, month=None, *args, **kwargs):
        """Get a file of a given type for a given kepid/quarter[/month]

        Inputs:
        ---------
        kepid   (int)
            Kepid of target of interest
        quarter (int)
            Quarter of interest
        isFluxFile (bool)
            If True, a lightcurve file is returned, otherwise a target
            pixel file is returned

        Optional Inputs:
        ----------------
        month   (int)
            If None, a long cadence file is returned, if a value
            of 1,2 or 3 is supplied, a short cadence file is returned.

        All other optional inputs are passed to ``pyfits.getdata()``

        Returns:
        -------------
        A FITs object, and possibly more, depending on the optional
        arguments to ``pyfits.getdata``
        """

        filename = self.getFilename(kepid, quarter, isFluxFile, month)
        localUrl = os.path.join( self.getLocalDir(quarter, kepid), filename)

        #Make remote url
        if isFluxFile:
            remotePath = self.remoteFluxPath
            compressedOnServer = False
        else:
            remotePath = self.remoteTpfPath
            compressedOnServer = True

        remoteUrl = self.makeRemoteUrl(remotePath, kepid, quarter, filename, compressedOnServer)
        return self.getData(localUrl, remoteUrl, compressedOnServer, *args, **kwargs)


    def makeRemoteUrl(self, remotePath, kepid, quarter, filename, compressedOnServer):
        """Figure out the URL where the desired file is stored.
        
        Returns a string, typically similar to
        http://archive.stsci.edu/$remotePath/.../$filename
        """
        raise NotImplementedError("Not Implemented in Base Class")

    def getFilename(self, kepid, quarter, isFluxFile, isShortCadence):
        """Figure out the name of the file you want to download
        
        (But not its path)
        """
        raise NotImplementedError("Not Implemented in Base Class")
Abstract classes for interfacing with the Mast Kepler Data archives.   Used by mastio.py Base class for queries of the MAST Kepler and K2 archives.     Handles submitting the URL, receiving and caching the results      The child class determines the url to submit Download data from MAST archive for a given filename and a         given quarter.          Inputs:         clobber:    If true, allow locally existing files to be overwritten         verbose:    If true, print the file meta-information #Write a zero length file to indicate no file is present Load data using favourite IO class                  Assumes that localUrl exists on disk. Default file format         for MAST is FITS, but sub-classes can override this method         for other file types as necessary Base class for classes interacting with Kepler/K2 data.          This class contains methods for querying for Lightcurve and      TPF files in either short or long cadence.          Example daughter classes are KeplerArchive and K2Archive.     Daughter classes must, at a minimum, reimplement getFilename()     and makeRemoteUrl(). You may want to write your own init, too. #Re-raising an exception makes the error easier to read #Re-raising an exception makes the error easier to read Get a long cadence lightcurve file.          A thin wrapper around getFile() Get a short cadence lightcurve file.          A thin wrapper around getFile() Get a long cadence TPF file.          A thin wrapper around getFile() Get a short cadence TPF file.          A thin wrapper around getFile() Get a file of a given type for a given kepid/quarter[/month]          Inputs:         ---------         kepid   (int)             Kepid of target of interest         quarter (int)             Quarter of interest         isFluxFile (bool)             If True, a lightcurve file is returned, otherwise a target             pixel file is returned          Optional Inputs:         ----------------         month   (int)             If None, a long cadence file is returned, if a value             of 1,2 or 3 is supplied, a short cadence file is returned.          All other optional inputs are passed to ``pyfits.getdata()``          Returns:         -------------         A FITs object, and possibly more, depending on the optional         arguments to ``pyfits.getdata`` #Make remote url Figure out the URL where the desired file is stored.                  Returns a string, typically similar to         http://archive.stsci.edu/$remotePath/.../$filename Figure out the name of the file you want to download                  (But not its path)