Source code for pyion.utils

""" 
# ===========================================================================
# Module with utility functions and classes for pyion
#
# Author: Marc Sanchez Net
# Date:   04/17/2019
# Copyright (c) 2019, California Institute of Technology ("Caltech").  
# U.S. Government sponsorship acknowledged.
# ===========================================================================
"""

# Generic imports
from abc import ABC
from datetime import datetime, timedelta
from enum import Enum, unique
from functools import wraps
import os
from pathlib import Path
import time

import pyion

# ============================================================================
# === Define variables/methods/classes that will be directly importable
# ============================================================================

__all__ = ['check_ion_env_vars']

# ============================================================================
# === Error messages
# ============================================================================

_env_var_missing = ('{} enviornment variable is not properly set',
                    'Current value is {}',
                    'You can either set it in bash using ``export ${}`` or',
                    'inside Python using ``pyion.{} = "..."``')

_path_missing = 'Path "{}" does not exist'

# ============================================================================
# === Abstract Proxy
# ============================================================================

class Proxy(ABC):
    """ Abstract proxy to ION. 

        Public Variables
        ----------------
        :ivar node_nbr:  Node number that this proxy is action upon.
        :ivar node_dir:  Directory for this node
        :ivar attached:  True if proxy is attached to ION.
    """
    def __init__(self, node_nbr):
        # Initialize variables
        self.node_nbr = node_nbr
        self.attached = False

        # Set the ION_NODE_LIST_DIR variable if necessary. This call resolves relative
        # paths, makes sure directory exits.
        set_ion_node_list_dir(pyion.ION_NODE_LIST_DIR)

        # Get ION_NODE_LIST_DIR environment variable
        nodes_path = os.environ.get('ION_NODE_LIST_DIR')

        # Set node number and node dir
        self.node_dir = Path(nodes_path)/str(node_nbr) if nodes_path else None

    def __str__(self):
        return '<{}: {} ({})>'.format(self.__class__.__name__, self.node_nbr,
                                      'Attached' if self.attached else 'Detached')

    def __repr__(self):
        return str(self)

# ============================================================================
# === Helper functions
# ============================================================================

def in_ion_folder(func):
    """ Decorator to change directory inside a node's folder to run ION
        commands. It can only be used from within the ``Proxy`` class
    """
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        # Get the node's directory
        node_dir = getattr(self, 'node_dir')

        # If no directory specified, just run. This is always the case
        # unless you run multiple ION nodes in a single machine.
        if node_dir is None: return func(self, *args, **kwargs)
    
        # Store current working directory
        cur_dir = os.getcwd()

        # Go to the node's directory
        os.chdir(str(node_dir.absolute()))

        # Execture function
        try:
            ret = func(self, *args, **kwargs)
        finally:
            # Go back to the previous working directory
            os.chdir(cur_dir)

        return ret
    return wrapper

def _chk_attached(func):
    """ Decorator that checks if Proxy is attached """
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        # You need to be attached
        if not self.attached:
            raise IOError('Not attached to ION, call ``attach`` first.')

        # Call function
        return func(self, *args, **kwargs)
    return wrapper

def _chk_is_open(func):
    """ Decorator that checks if PoP is opened """
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        # You need to be attached
        if not self.is_open:
            raise ConnectionAbortedError('{}'.format(self))

        # Call function
        return func(self, *args, **kwargs)
    return wrapper

[docs] def check_ion_env_vars(ION_NODE_LIST_DIR): """ Check the ION environment variables to ensure they are consistent (e.g., the paths set are valid and exist in the host). Host environment variables are only utilized if the module variables ``pyion.ION_PWD`` and ``pyion.ION_NODE_LIST_DIR`` are None. """ # Initialize variables ion_nld = ION_NODE_LIST_DIR # If not set, use environment variables if ion_nld is None: ion_nld = os.environ.get('ION_NODE_LIST_DIR') # Convert to valid path try: ion_nld = Path(ion_nld).resolve().absolute() except TypeError: msg = '\n'.join(_env_var_missing) raise OSError(msg.format('ION_NODE_LIST_DIR', ion_nld, 'ION_NODE_LIST_DIR', 'ION_NODE_LIST_DIR')) # If path does not exists, raise error if not ion_nld.exists(): raise IOError(_path_missing.format(ion_nld)) # Set the module variables return ion_nld
def set_ion_node_list_dir(path): """ Set the ION_NODE_LIST_DIR environment variable. Is not a valid OS path, an exception is raise. :param str: Relative or absolute path. """ # If empty, just return if path is None: return # Check consistency of this variable ION_NODE_LIST_DIR = check_ion_env_vars(path) # Set it in the environment os.environ['ION_NODE_LIST_DIR'] = str(ION_NODE_LIST_DIR) def _register_proxy(proxy_map, key, proxy_cls, *args, **kwargs): """ Implements the singleton mechanism for proxies """ # If it already exists, return it if key in proxy_map: return proxy_map[key] # Create a new one. It gets registered during initialization proxy = proxy_cls(*args, **kwargs) # Register this proxy proxy_map[key] = proxy # Return new proxy return proxy def _unregister_proxy(proxy_map, node_nbr): # Try to unregister node. If not possible, return try: del proxy_map[str(node_nbr)] except KeyError: pass def rel2abs_time(str_time): """ Assumes that relative times are provided as ``+xxxx`` where xxxx is the time gap in seconds. """ # If this is an absolute time, return if '+' not in str_time: return str_time # Compute the date from now tstart = datetime.now() + timedelta(seconds=int(str_time[1:])) # Transform to ION time format return tstart.strftime('%Y/%m/%d-%H:%M:%S') class Rate(): def __init__(self, rate): self.rate = rate self.tic = time.time() def sleep(self): # Sleep for a while time.sleep(max(0, 1.0/self.rate - (time.time()-self.tic))) # Reset tic self.tic = time.time()