<commit_before>from __future__ import print_function, division
import numpy as np
from timeit import default_timer as timer
from pyscf.nao.tddft_iter import use_numba, tddft_iter
from pyscf.nao import gw
from scipy.linalg import blas
from pyscf.nao.m_pack2den import pack2den_u, pack2den_l
if use_numba: from pyscf.nao.m_iter_div_eigenenergy_numba import div_eigenenergy_numba

class bse_iter(gw):

  def __init__(self, **kw):
    """
      Iterative BSE a la PK, DF, OC JCTC
      additionally to the fields from tddft_iter_c, we add the dipole matrix elements dab[ixyz][a,b]
      which is constructed as list of numpy arrays
       $ d_i = \int f^a(r) r_i f^b(r) dr $
    """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    xc_code_kw = kw['xc_code'] if 'xc_code' in kw else None
    gw.__init__(self, **kw)
    self.l0_ncalls = 0
    self.dip_ab = [d.toarray() for d in self.dipole_coo()]
    self.norbs2 = self.norbs**2
    kernel_den = pack2den_l(self.kernel)
    n = self.norbs
    v_dab = self.v_dab
    cc_da = self.cc_da
    self.x = self.mo_coeff[0,0,:,:,0]
    self.kernel_4p = (((v_dab.T*(cc_da*kernel_den))*cc_da.T)*v_dab).reshape([n*n,n*n])
    #print(type(self.kernel_4p), self.kernel_4p.shape, 'this is just a reference kernel, must be removed later for sure')

    self.xc_code = self.xc_code if xc_code_kw is None else xc_code_kw
    xc = self.xc_code.split(',')[0]
    if self.verbosity>0:
      print(__name__, '     xc_code_mf ', self.xc_code_mf)
      print(__name__, ' xc_code_kernel ', self.xc_code_kernel)
      print(__name__, '    xc_code_scf ', self.xc_code_scf)
      print(__name__, '        xc_code ', self.xc_code)

    if xc=='CIS' or xc=='HF' or xc=='GW':
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', self.kernel_4p.reshape([n,n,n,n])).reshape([n*n,n*n])
    elif xc=='RPA' or xc=='LDA' or xc=='GGA':
      pass
    else :
      print(' ?? xc_code ', self.xc_code, xc)
      raise RuntimeError('??? xc_code ???')

    if xc=='GW':
      w_c_4p = (((v_dab.T*(cc_da*  self.si_c([0.0])[0].real ))*cc_da.T)*v_dab).reshape([n*n,n*n])
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', w_c_4p.reshape([n,n,n,n])).reshape([n*n,n*n])

    if hasattr(self, 'mo_energy_gw'):
      # Need to redefine a lot...
      self.mo_energy_scf = np.copy(self.mo_energy)
      self.mo_energy = self.mo_energy_gw
      self.mo_coeff_scf = np.copy(self.mo_coeff)
      self.mo_coeff = self.mo_coeff_gw
      self.x = self.mo_coeff[0,0,:,:,0]
      self.ksn2e = np.require(np.zeros((1,self.nspin,self.norbs)), dtype=self.dtype, requirements='CW')
      self.ksn2e[0,0,:] = self.mo_energy
      ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
      if all(ksn2fd[0,0,:]>self.nfermi_tol):
        print(__name__, self.telec, nfermi_tol, ksn2fd[0,0,:])
        raise RuntimeError('telec is too high?')

      self.ksn2f = (3-self.nspin)*ksn2fd
      self.nfermi = np.argmax(ksn2fd[0,0,:]<self.nfermi_tol)
      self.vstart = np.argmax(1.0-ksn2fd[0,0,:]>=self.nfermi_tol)
      self.xocc = self.mo_coeff[0,0,0:self.nfermi,:,0]
      self.xvrt = self.mo_coeff[0,0,self.vstart:,:,0]

  def apply_l0(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nm2v[self.vstart:, :self.nfermi] = blas.cgemm(1.0, np.dot(self.xvrt, sab), self.xocc.T)
    nm2v[:self.nfermi, self.vstart:] = blas.cgemm(1.0, np.dot(self.xocc, sab), self.xvrt.T)

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_exp(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1
    nb2v = np.dot(self.x, sab)
    nm2v = blas.cgemm(1.0, nb2v, self.x.T)
    print(nm2v.dtype)
    print(nm2v[self.vstart:, :self.nfermi])
    print(nm2v[:self.nfermi, self.vstart:])

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nb2v1 = np.dot(self.xocc, sab)
    nm2v1 = blas.cgemm(1.0, nb2v1, self.xvrt.T)

    nb2v2 = np.dot(self.xvrt, sab)
    nm2v2 = blas.cgemm(1.0, nb2v2, self.xocc.T)

    nm2v[self.vstart:, :self.nfermi] = nm2v2
    nm2v[:self.nfermi, self.vstart:] = nm2v1

    print(nm2v.dtype, nm2v1.shape, nm2v1.dtype)
    print(nm2v[self.vstart:, :self.nfermi])
    print(nm2v[:self.nfermi, self.vstart:])
    #raise RuntimeError('11')

    #nm2v2 = np.copy(nm2v1)
    #for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:self.nfermi],self.ksn2f[0,0,:self.nfermi])):
      #for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,self.vstart:],self.ksn2f[0,0,self.vstart:])):
        #nm2v1[n,m] = nm2v1[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    #for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:self.nfermi],self.ksn2f[0,0,:self.nfermi])):
      #for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,self.vstart:],self.ksn2f[0,0,self.vstart:])):
        #nm2v2[n,m] = nm2v2[n,m] * (fm-fn) * ( 1.0 / (comega - (en - em)))

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_ref(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1
    nb2v = np.dot(self.x, sab)
    nm2v = blas.cgemm(1.0, nb2v, self.x.T)

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def seff(self, sext, comega=1j*0.0):
    """ This computes an effective two point field (scalar non-local potential) given an external two point field.
        L = L0 (1 - K L0)^-1
        We want therefore an effective X_eff for a given X_ext
        X_eff = (1 - K L0)^-1 X_ext   or   we need to solve linear equation
        (1 - K L0) X_eff = X_ext

        The operator (1 - K L0) is named self.sext2seff_matvec """

    from scipy.sparse.linalg import gmres, lgmres as gmres_alias, LinearOperator
    assert sext.size==(self.norbs2), "%r,%r"%(sext.size,self.norbs2)

    self.comega_current = comega
    op = LinearOperator((self.norbs2,self.norbs2), matvec=self.sext2seff_matvec, dtype=self.dtypeComplex)
    sext_shape = np.require(sext.reshape(self.norbs2), dtype=self.dtypeComplex, requirements='C')
    resgm,info = gmres_alias(op, sext_shape, tol=self.tddft_iter_tol)
    return (resgm.reshape([self.norbs,self.norbs]),info)

  def sext2seff_matvec(self, sab):
    """ This is operator which we effectively want to have inverted (1 - K L0) and find the action of it's
    inverse by solving a linear equation with a GMRES method. See the method seff(...)"""
    self.matvec_ncalls+=1

    l0 = self.apply_l0(sab, self.comega_current).reshape(self.norbs2)

    l0_reim = np.require(l0.real, dtype=self.dtype, requirements=["A", "O"])     # real part
    mv_real = np.dot(self.kernel_4p, l0_reim)

    l0_reim = np.require(l0.imag, dtype=self.dtype, requirements=["A", "O"])     # imaginary part
    mv_imag = np.dot(self.kernel_4p, l0_reim)

    return sab - (mv_real + 1.0j*mv_imag)

  def apply_l(self, sab, comega=1j*0.0):
    """ This applies the interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    seff,info = self.seff(sab, comega)
    return self.apply_l0( seff, comega )

  def comp_polariz_nonin_ave(self, comegas):
    """ Non-interacting average polarizability """
    p = np.zeros(len(comegas), dtype=self.dtypeComplex)
    for ixyz in range(3):
      for iw,omega in enumerate(comegas):
        vab = self.apply_l0(self.dip_ab[ixyz], omega)
        p[iw] += (vab*self.dip_ab[ixyz]).sum()/3.0
    return p

  def comp_polariz_inter_ave(self, comegas):
    """ Compute a direction-averaged interacting polarizability  """
    p = np.zeros(len(comegas), dtype=self.dtypeComplex)
    for ixyz in range(3):
      for iw,omega in enumerate(comegas):
        if self.verbosity>1: print(__name__, ixyz, iw)
        vab = self.apply_l(self.dip_ab[ixyz], omega)
        p[iw] += (vab*self.dip_ab[ixyz]).sum()/3.0
    return p
<commit_msg>easier to not redefine mo_energy<commit_after>from __future__ import print_function, division
import numpy as np
from timeit import default_timer as timer
from pyscf.nao.tddft_iter import use_numba, tddft_iter
from pyscf.nao import gw
from scipy.linalg import blas
from pyscf.nao.m_pack2den import pack2den_u, pack2den_l
if use_numba: from pyscf.nao.m_iter_div_eigenenergy_numba import div_eigenenergy_numba

class bse_iter(gw):

  def __init__(self, **kw):
    """
      Iterative BSE a la PK, DF, OC JCTC
      additionally to the fields from tddft_iter_c, we add the dipole matrix elements dab[ixyz][a,b]
      which is constructed as list of numpy arrays
       $ d_i = \int f^a(r) r_i f^b(r) dr $
    """
    from pyscf.nao.m_fermi_dirac import fermi_dirac_occupations

    xc_code_kw = kw['xc_code'] if 'xc_code' in kw else None
    gw.__init__(self, **kw)
    self.l0_ncalls = 0
    self.dip_ab = [d.toarray() for d in self.dipole_coo()]
    self.norbs2 = self.norbs**2
    kernel_den = pack2den_l(self.kernel)
    n = self.norbs
    v_dab = self.v_dab
    cc_da = self.cc_da
    self.x = self.mo_coeff[0,0,:,:,0]
    self.kernel_4p = (((v_dab.T*(cc_da*kernel_den))*cc_da.T)*v_dab).reshape([n*n,n*n])
    #print(type(self.kernel_4p), self.kernel_4p.shape, 'this is just a reference kernel, must be removed later for sure')

    self.xc_code = self.xc_code if xc_code_kw is None else xc_code_kw
    xc = self.xc_code.split(',')[0]
    if self.verbosity>0:
      print(__name__, '     xc_code_mf ', self.xc_code_mf)
      print(__name__, ' xc_code_kernel ', self.xc_code_kernel)
      print(__name__, '    xc_code_scf ', self.xc_code_scf)
      print(__name__, '        xc_code ', self.xc_code)

    if xc=='CIS' or xc=='HF' or xc=='GW':
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', self.kernel_4p.reshape([n,n,n,n])).reshape([n*n,n*n])
    elif xc=='RPA' or xc=='LDA' or xc=='GGA':
      pass
    else :
      print(' ?? xc_code ', self.xc_code, xc)
      raise RuntimeError('??? xc_code ???')

    if xc=='GW':
      w_c_4p = (((v_dab.T*(cc_da*  self.si_c([0.0])[0].real ))*cc_da.T)*v_dab).reshape([n*n,n*n])
      self.kernel_4p -= 0.5*np.einsum('abcd->bcad', w_c_4p.reshape([n,n,n,n])).reshape([n*n,n*n])

    if hasattr(self, 'mo_energy_gw') and xc=='GW':  # Need to redefine eigenvectors previously set in the class tddft_iter
      self.ksn2e = np.require(np.zeros((1,self.nspin,self.norbs)), dtype=self.dtype, requirements='CW')
      self.ksn2e[0,0,:] = self.mo_energy_gw

      ksn2fd = fermi_dirac_occupations(self.telec, self.ksn2e, self.fermi_energy)
      if all(ksn2fd[0,0,:]>self.nfermi_tol):
        print(__name__, self.telec, nfermi_tol, ksn2fd[0,0,:])
        raise RuntimeError('telec is too high?')
      self.ksn2f = (3-self.nspin)*ksn2fd

      self.x = self.mo_coeff_gw[0,0,:,:,0]
      self.nfermi = np.argmax(ksn2fd[0,0,:]<self.nfermi_tol)
      self.vstart = np.argmax(1.0-ksn2fd[0,0,:]>=self.nfermi_tol)
      self.xocc = self.mo_coeff_gw[0,0,0:self.nfermi,:,0]
      self.xvrt = self.mo_coeff_gw[0,0,self.vstart:,:,0]

  def apply_l0(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nm2v[self.vstart:, :self.nfermi] = blas.cgemm(1.0, np.dot(self.xvrt, sab), self.xocc.T)
    nm2v[:self.nfermi, self.vstart:] = blas.cgemm(1.0, np.dot(self.xocc, sab), self.xvrt.T)

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_exp(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1
    nb2v = np.dot(self.x, sab)
    nm2v = blas.cgemm(1.0, nb2v, self.x.T)
    print(nm2v.dtype)
    print(nm2v[self.vstart:, :self.nfermi])
    print(nm2v[:self.nfermi, self.vstart:])

    nm2v = np.zeros((self.norbs,self.norbs), self.dtypeComplex)
    nb2v1 = np.dot(self.xocc, sab)
    nm2v1 = blas.cgemm(1.0, nb2v1, self.xvrt.T)

    nb2v2 = np.dot(self.xvrt, sab)
    nm2v2 = blas.cgemm(1.0, nb2v2, self.xocc.T)

    nm2v[self.vstart:, :self.nfermi] = nm2v2
    nm2v[:self.nfermi, self.vstart:] = nm2v1

    print(nm2v.dtype, nm2v1.shape, nm2v1.dtype)
    print(nm2v[self.vstart:, :self.nfermi])
    print(nm2v[:self.nfermi, self.vstart:])
    #raise RuntimeError('11')

    #nm2v2 = np.copy(nm2v1)
    #for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:self.nfermi],self.ksn2f[0,0,:self.nfermi])):
      #for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,self.vstart:],self.ksn2f[0,0,self.vstart:])):
        #nm2v1[n,m] = nm2v1[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    #for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:self.nfermi],self.ksn2f[0,0,:self.nfermi])):
      #for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,self.vstart:],self.ksn2f[0,0,self.vstart:])):
        #nm2v2[n,m] = nm2v2[n,m] * (fm-fn) * ( 1.0 / (comega - (en - em)))

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def apply_l0_ref(self, sab, comega=1j*0.0):
    """ This applies the non-interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    assert sab.size==(self.norbs2), "%r,%r"%(sab.size,self.norbs2)

    sab = sab.reshape([self.norbs,self.norbs])
    self.l0_ncalls+=1
    nb2v = np.dot(self.x, sab)
    nm2v = blas.cgemm(1.0, nb2v, self.x.T)

    for n,[en,fn] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
      for m,[em,fm] in enumerate(zip(self.ksn2e[0,0,:],self.ksn2f[0,0,:])):
        nm2v[n,m] = nm2v[n,m] * (fn-fm) * ( 1.0 / (comega - (em - en)))

    nb2v = blas.cgemm(1.0, nm2v, self.x)
    ab2v = blas.cgemm(1.0, self.x.T, nb2v)
    return ab2v

  def seff(self, sext, comega=1j*0.0):
    """ This computes an effective two point field (scalar non-local potential) given an external two point field.
        L = L0 (1 - K L0)^-1
        We want therefore an effective X_eff for a given X_ext
        X_eff = (1 - K L0)^-1 X_ext   or   we need to solve linear equation
        (1 - K L0) X_eff = X_ext

        The operator (1 - K L0) is named self.sext2seff_matvec """

    from scipy.sparse.linalg import gmres, lgmres as gmres_alias, LinearOperator
    assert sext.size==(self.norbs2), "%r,%r"%(sext.size,self.norbs2)

    self.comega_current = comega
    op = LinearOperator((self.norbs2,self.norbs2), matvec=self.sext2seff_matvec, dtype=self.dtypeComplex)
    sext_shape = np.require(sext.reshape(self.norbs2), dtype=self.dtypeComplex, requirements='C')
    resgm,info = gmres_alias(op, sext_shape, tol=self.tddft_iter_tol)
    return (resgm.reshape([self.norbs,self.norbs]),info)

  def sext2seff_matvec(self, sab):
    """ This is operator which we effectively want to have inverted (1 - K L0) and find the action of it's
    inverse by solving a linear equation with a GMRES method. See the method seff(...)"""
    self.matvec_ncalls+=1

    l0 = self.apply_l0(sab, self.comega_current).reshape(self.norbs2)

    l0_reim = np.require(l0.real, dtype=self.dtype, requirements=["A", "O"])     # real part
    mv_real = np.dot(self.kernel_4p, l0_reim)

    l0_reim = np.require(l0.imag, dtype=self.dtype, requirements=["A", "O"])     # imaginary part
    mv_imag = np.dot(self.kernel_4p, l0_reim)

    return sab - (mv_real + 1.0j*mv_imag)

  def apply_l(self, sab, comega=1j*0.0):
    """ This applies the interacting four point Green's function to a suitable vector (e.g. dipole matrix elements)"""
    seff,info = self.seff(sab, comega)
    return self.apply_l0( seff, comega )

  def comp_polariz_nonin_ave(self, comegas):
    """ Non-interacting average polarizability """
    p = np.zeros(len(comegas), dtype=self.dtypeComplex)
    for ixyz in range(3):
      for iw,omega in enumerate(comegas):
        vab = self.apply_l0(self.dip_ab[ixyz], omega)
        p[iw] += (vab*self.dip_ab[ixyz]).sum()/3.0
    return p

  def comp_polariz_inter_ave(self, comegas):
    """ Compute a direction-averaged interacting polarizability  """
    p = np.zeros(len(comegas), dtype=self.dtypeComplex)
    for ixyz in range(3):
      for iw,omega in enumerate(comegas):
        if self.verbosity>1: print(__name__, ixyz, iw)
        vab = self.apply_l(self.dip_ab[ixyz], omega)
        p[iw] += (vab*self.dip_ab[ixyz]).sum()/3.0
    return p
<|endoftext|>


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

import argparse
import subprocess
import os
import logging
import zipfile
from datetime import datetime
from progress.bar import Bar

class Dinghy:

    def __init__(self, system_directories, show_progress):
        # Refuse to do anything if RunDeck is running
        # This is best practice according to the docs:
        # http://rundeck.org/2.6.11/administration/backup-and-recovery.html
        if self._rundeck_is_running():
            logging.error("rundeckd cannot be running while you take a backup"
                          " or restore from backup")
            raise Exception("rundeckd is still running")

        self.count = 0
        self.bar = None
        self.show_progress = show_progress
        # Directories to include in backup and restore
        self.system_directories = system_directories

    def _rundeck_is_running(self):
        """
        Returns True if rundeckd is running, False otherwise
        """
        try:
            status = subprocess.check_output(
                ["service", "rundeckd", "status"],
                # Universal newlines ensures error.output is a string
                universal_newlines=True
            )
        except subprocess.CalledProcessError as error:
            if "rundeckd is not running" not in error.output:
                raise Exception("error running service command")
            else:
                status = error.output
        if "rundeckd is running" in status:
            return True
        else:
            return False

    def _add_directory_to_zip(self, zip_handle, path):

        for root, dirs, files in os.walk(path):
            logging.debug("directory: {}".format(root))
            for f in files:
                if self.show_progress:
                    self.bar.next()
                logging.debug("file: {}".format(f))
                zip_handle.write(os.path.join(root, f))

    def backup(self, destination_path, filename):
        # Start message
        logging.debug("starting backup")

        # Fail if backup dir is not found
        if not os.path.exists(destination_path):
            logging.error("backup directory {} not found".format(
                destination_path))
            raise Exception("could not find backup directory")

        file_path = os.path.join(destination_path, filename)
        logging.debug("using full backup path {}".format(file_path))

        if self.show_progress:
            # Count files in all directories
            logging.info("counting files...")
            for directory in self.system_directories:
                for root, dirs, files in os.walk(directory):
                    for f in files:
                        self.count += 1
            logging.debug("total number of files is {}".format(self.count))
            # Create progress bar
            self.bar = Bar('Processing', max=self.count)

        # Create zip file and save all directories to it
        # allowZip64 must be True to allow file size > 4GB
        with zipfile.ZipFile(file_path, 'w', allowZip64=True) as zip_file:
            for directory in self.system_directories:
                logging.info("adding directory {}".format(directory))
                self._add_directory_to_zip(zip_file, directory)
                # To get the progress bar on separate line from
                # log messages when printing log to stdout
                print("")

        if self.show_progress:
            # Close progress bar
            self.bar.finish()

    def restore(self, filename):
        return

def main(arguments):
    # Gather arguments
    parser_name = arguments.subparser_name
    debug_mode = arguments.debug
    if arguments.no_progress:
        show_progress = False
    else:
        show_progress = True

    # Enable debug logging if flag is set
    if debug_mode:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO
    # Set up logging
    logging.basicConfig(
        format='%(asctime)s %(levelname)s %(message)s',
        level=log_level)

    # Set backup directories
    if arguments.dirs:
        print(arguments.dirs)
        system_directories = arguments.dirs[0].split(",")
        # Validate that no space is left in input
        for directory in system_directories:
            if os.path.exists(directory) or os.access(directory, os.W_OK):
                pass
            else:
                raise Exception("directory is not a valid path "
                                "or not writeable: {}".format(directory))
    else:
        # Default is to backup and restore all directories
        system_directories = [
            "/var/lib/rundeck/data",          # database
            "/var/lib/rundeck/logs",          # execution logs (by far biggest)
            "/var/lib/rundeck/.ssh",          # ssh keys
            "/var/lib/rundeck/var/storage",   # key storage files and metadata
            "/var/rundeck/projects"           # project definitions
        ]

    dinghy = Dinghy(system_directories=system_directories,
                    show_progress=show_progress)

    if parser_name == "backup":
        # Run backup
        dinghy.backup(destination_path=arguments.dest,
                      filename=arguments.filename)
    elif parser_name == "restore":
        dinghy.restore(filename=arguments.file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='dinghy: helper for backup and restore of RunDeck')
    parser.add_argument(
        '--debug',
        '-d',
        action='store_true',
        help='enable debug logging')
    parser.add_argument(
        '--no-progress',
        action='store_true',
        help='disable progress bar')
    parser.add_argument(
        '--dirs',
        type=str,
        nargs="*",
        help='comma-separated list that overrides the default list of system'
             'directories to backup/restore')

    subparsers = parser.add_subparsers(help='command help',
                                       dest='subparser_name')

    # Backup options
    backup_parser = subparsers.add_parser('backup', help='create a backup')
    backup_parser.add_argument(
        '--dest',
        type=str,
        required=True,
        help='path to write backup file to')
    backup_parser.add_argument(
        '--filename',
        type=str,
        help='override the filename used the for backup file',
        default='rundeck-backup-{}.zip'.format(
            datetime.now().strftime('%Y-%M-%d--%H-%m-%S')))

    # Restore options
    restore_parser = subparsers.add_parser(
        'restore',
        help='restore from a backup file')
    restore_parser.add_argument(
        '--file',
        type=str,
        help='path to backup file to restore from')
    main(parser.parse_args())
<commit_msg>added check for conflicting paths<commit_after>#!/usr/bin/env python

import argparse
import subprocess
import os
import logging
import zipfile
from datetime import datetime
from progress.bar import Bar

class Dinghy:

    def __init__(self, system_directories, show_progress):
        # Refuse to do anything if RunDeck is running
        # This is best practice according to the docs:
        # http://rundeck.org/2.6.11/administration/backup-and-recovery.html
        if self._rundeck_is_running():
            logging.error("rundeckd cannot be running while you take a backup"
                          " or restore from backup")
            raise Exception("rundeckd is still running")

        self.count = 0
        self.bar = None
        self.show_progress = show_progress
        # Directories to include in backup and restore
        self.system_directories = system_directories

        # Raise exception if duplicate or
        # overlapping directories are passed in
        if self._has_duplicate_or_overlap(self.system_directories):
            raise Exception("duplicate or overlapping directories detected")

    def _has_duplicate_or_overlap(self, paths):
        """
        Returns true if there are duplicate or overlapping
        paths in a list of paths
        """
        if len(paths) > 1:
            first = paths[0]
            remaining = paths[1:]
            for item in remaining:
                if first in item or item in first:
                    logging.error("found conflicting paths {},{}".format(
                        first,
                        item
                    ))
                    return True
            self._has_duplicate_or_overlap(remaining)
        return False

    def _rundeck_is_running(self):
        """
        Returns True if rundeckd is running, False otherwise
        """
        try:
            status = subprocess.check_output(
                ["service", "rundeckd", "status"],
                # Universal newlines ensures error.output is a string
                universal_newlines=True
            )
        except subprocess.CalledProcessError as error:
            if "rundeckd is not running" not in error.output:
                raise Exception("error running service command")
            else:
                status = error.output
        if "rundeckd is running" in status:
            return True
        else:
            return False

    def _add_directory_to_zip(self, zip_handle, path):
        for root, dirs, files in os.walk(path):
            logging.debug("directory: {}".format(root))
            for f in files:
                logging.debug("file: {}".format(f))
                zip_handle.write(os.path.join(root, f))
                if self.show_progress:
                    self.bar.next()

    def backup(self, destination_path, filename):
        # Start message
        logging.debug("starting backup")

        # Fail if backup dir is not found
        if not os.path.exists(destination_path):
            logging.error("backup directory {} not found".format(
                destination_path))
            raise Exception("could not find backup directory")

        file_path = os.path.join(destination_path, filename)
        logging.debug("using full backup path {}".format(file_path))

        if self.show_progress:
            # Count files in all directories
            logging.info("counting files...")
            for directory in self.system_directories:
                for root, dirs, files in os.walk(directory):
                    for f in files:
                        self.count += 1
            logging.debug("total number of files is {}".format(self.count))
            # Create progress bar
            self.bar = Bar('Processing', max=self.count)

        # Create zip file and save all directories to it
        # allowZip64 must be True to allow file size > 4GB
        with zipfile.ZipFile(file_path, 'w', allowZip64=True) as zip_file:
            for directory in self.system_directories:
                logging.info("adding directory {}".format(directory))
                self._add_directory_to_zip(zip_file, directory)
                # To get the progress bar on separate line from
                # log messages when printing log to stdout
                print("")

        if self.show_progress:
            # Close progress bar
            self.bar.finish()

    def restore(self, filename):
        return

def main(arguments):
    # Gather arguments
    parser_name = arguments.subparser_name
    debug_mode = arguments.debug
    if arguments.no_progress:
        show_progress = False
    else:
        show_progress = True

    # Enable debug logging if flag is set
    if debug_mode:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO
    # Set up logging
    logging.basicConfig(
        format='%(asctime)s %(levelname)s %(message)s',
        level=log_level)

    # Set backup directories
    if arguments.dirs:
        print(arguments.dirs)
        system_directories = arguments.dirs[0].split(",")
        # Validate that no space is left in input
        for directory in system_directories:
            if os.path.exists(directory) or os.access(directory, os.W_OK):
                pass
            else:
                raise Exception("directory is not a valid path "
                                "or not writeable: {}".format(directory))
        logging.warning("overriding default directories with {}".format(
            system_directories
        ))
    else:
        # Default is to backup and restore all directories
        system_directories = [
            "/var/lib/rundeck/data",          # database
            "/var/lib/rundeck/logs",          # execution logs (by far biggest)
            "/var/lib/rundeck/.ssh",          # ssh keys
            "/var/lib/rundeck/var/storage",   # key storage files and metadata
            "/var/rundeck/projects"           # project definitions
        ]

    dinghy = Dinghy(system_directories=system_directories,
                    show_progress=show_progress)

    if parser_name == "backup":
        # Run backup
        dinghy.backup(destination_path=arguments.dest,
                      filename=arguments.filename)
    elif parser_name == "restore":
        dinghy.restore(filename=arguments.file)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='dinghy: helper for backup and restore of RunDeck')
    parser.add_argument(
        '--debug',
        '-d',
        action='store_true',
        help='enable debug logging')
    parser.add_argument(
        '--no-progress',
        action='store_true',
        help='disable progress bar')
    parser.add_argument(
        '--dirs',
        type=str,
        nargs="*",
        help='comma-separated list that overrides the default list of system'
             'directories to backup/restore')

    subparsers = parser.add_subparsers(help='command help',
                                       dest='subparser_name')

    # Backup options
    backup_parser = subparsers.add_parser('backup', help='create a backup')
    backup_parser.add_argument(
        '--dest',
        type=str,
        required=True,
        help='path to write backup file to')
    backup_parser.add_argument(
        '--filename',
        type=str,
        help='override the filename used the for backup file',
        default='rundeck-backup-{}.zip'.format(
            datetime.now().strftime('%Y-%M-%d--%H-%m-%S')))

    # Restore options
    restore_parser = subparsers.add_parser(
        'restore',
        help='restore from a backup file')
    restore_parser.add_argument(
        '--file',
        type=str,
        help='path to backup file to restore from')
    main(parser.parse_args())
<|endoftext|>


<commit_before># Copyright 2019 Ram Rachum and collaborators.
# This program is distributed under the MIT license.
'''
PySnooper - Never use print for debugging again

Usage:

    import pysnooper

    @pysnooper.snoop()
    def your_function(x):
        ...

A log will be written to stderr showing the lines executed and variables
changed in the decorated function.

For more information, see https://github.com/cool-RR/PySnooper
'''

from .tracer import Tracer as snoop
from .variables import Attrs, Exploding, Indices, Keys
import collections

__VersionInfo = collections.namedtuple('VersionInfo',
                                       ('major', 'minor', 'micro'))

__version__ = '0.0.35'
__version_info__ = __VersionInfo(*(map(int, __version__.split('.'))))

del collections, __VersionInfo # Avoid polluting the namespace
<commit_msg>Bump version<commit_after># Copyright 2019 Ram Rachum and collaborators.
# This program is distributed under the MIT license.
'''
PySnooper - Never use print for debugging again

Usage:

    import pysnooper

    @pysnooper.snoop()
    def your_function(x):
        ...

A log will be written to stderr showing the lines executed and variables
changed in the decorated function.

For more information, see https://github.com/cool-RR/PySnooper
'''

from .tracer import Tracer as snoop
from .variables import Attrs, Exploding, Indices, Keys
import collections

__VersionInfo = collections.namedtuple('VersionInfo',
                                       ('major', 'minor', 'micro'))

__version__ = '0.0.36'
__version_info__ = __VersionInfo(*(map(int, __version__.split('.'))))

del collections, __VersionInfo # Avoid polluting the namespace
<|endoftext|>


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

import sys
import time
import uuid
import select
import os
import json
import socket
import shutil
import errno
from SocketServer import (ThreadingUnixStreamServer,
                          StreamRequestHandler)
from io import BytesIO

from pysource import env
from pysource import ExecutionError
from pysource import remote_call_handlers
from pysource import request_context

RESPONSE_STATUS_OK = "ok"
RESPONSE_STATUS_ERROR = "error"

CONTROL_SOCKET_ERROR_CODE = 1
CONTROL_SOCKET_LENGTH_CODE = 2

unix_socket_path = os.path.join(env.pysource_dir, 'socket')

def _handle(req_type, payload, **kwargs):
    handler = remote_call_handlers.get(req_type)
    if handler is None:
        raise RuntimeError('Unknown request type: {0}'.format(req_type))
    return handler(**payload)

class RequestHandler(StreamRequestHandler):

    def __init__(self, request, client_address, server):
        self.pipe_control_handler = None
        StreamRequestHandler.__init__(self, request, client_address, server)

    def handle(self):
        try:
            res_status = RESPONSE_STATUS_OK
            body = _read_body(self.rfile)
            if body['piped'] is True:
                request_context.piped = True
                self.pipe_control_handler = \
                    PipeControlSocketHandler(body['uid'])
                self._handle_piped(body, self.pipe_control_handler)
                res_payload = {}
            else:
                res_payload = _handle(**body)
        except Exception, e:
            res_status = RESPONSE_STATUS_ERROR
            res_payload = {'error': str(e)}

        if self.pipe_control_handler:
            response_file = self.pipe_control_handler.wfile
            if res_status == RESPONSE_STATUS_ERROR:
                try:
                    self.wfile.flush()
                except socket.error:
                    pass
                response_file.write('{}\r\n'.format(CONTROL_SOCKET_ERROR_CODE))
                response_file.flush()
        else:
            response_file = self.wfile
        _write_body(response_file, {
            'payload': res_payload,
            'status': res_status
        })

    def finish(self):
        StreamRequestHandler.finish(self)
        if self.pipe_control_handler:
            self.pipe_control_handler.close()

    def _handle_piped(self, body, pipe_control_handler):
        self.connection.setblocking(0)
        try:
            pipe_control_handler.accept()
            piped_input = PipeControlledInputSocket(
                self.rfile,
                self.connection,
                pipe_control_handler.rfile,
                pipe_control_handler.conn)
            piped_output = PipeControlledOutputSocket(
                self.wfile,
                self.connection,
                pipe_control_handler.wfile,
                pipe_control_handler.conn)
            request_context.stdin = piped_input
            request_context.stdout = piped_output
            result = _handle(**body)
            if result:
                piped_output.write(str(result))
            piped_input.close()
            piped_output.close()
        finally:
            self.connection.setblocking(1)
            pipe_control_handler.conn.setblocking(1)

class PipeControlSocketHandler(object):

    def __init__(self, uid):
        self.uid = uid
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.unix_socket_path = os.path.join(env.pysource_dir, self.uid)
        self.conn = None
        self.wfile = None
        self.rfile = None
        self.server = True

    def accept(self):
        self.sock.bind(self.unix_socket_path)
        self.sock.listen(1)
        conn, addr = self.sock.accept()
        self._setup(conn)

    def connect(self):
        self.server = False
        retries = 3
        current_tries = 0
        while True:
            try:
                current_tries += 1
                self.sock.connect(self.unix_socket_path)
                break
            except socket.error, e:
                if e.errno == errno.ENOENT and current_tries < retries:
                    time.sleep(0.1)
                else:
                    raise

        self._setup(self.sock)

    def _setup(self, conn):
        self.conn = conn
        self.conn.setblocking(0)
        self.rfile = self.conn.makefile('rb', -1)
        self.wfile = self.conn.makefile('wb', 0)

    def close(self):
        if not self.wfile.closed:
            try:
                self.conn.setblocking(1)
                self.wfile.flush()
            except socket.error:
                pass
        self.wfile.close()
        self.rfile.close()
        if self.server:
            try:
                self.conn.shutdown(socket.SHUT_WR)
            except socket.error:
                pass
            try:
                self.sock.close()
            except socket.error:
                pass
        self.conn.close()
        if self.server and os.path.exists(self.unix_socket_path):
            os.remove(self.unix_socket_path)

class PipeControlledBaseSocket(object):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        self.closed = False
        self.data_file = data_file
        self.data_socket = data_socket
        self.control_file = control_file
        self.control_socket = control_socket
        self.sockets = [self.data_socket, self.control_socket]

class PipeControlledInputSocket(PipeControlledBaseSocket):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        super(PipeControlledInputSocket, self).__init__(data_file,
                                                        data_socket,
                                                        control_file,
                                                        control_socket)
        self.bytes_read = 0
        self.total_bytes = None
        self.done = False

    def read(self, length=0):
        if self.done:
            return None
        result = BytesIO()
        while True:
            if self.bytes_read == self.total_bytes:
                break
            readable, _, _ = select.select(self.sockets, [], [], 0.5)
            if self.control_socket in readable:
                control_code = int(self.control_file.readline())
                if control_code == CONTROL_SOCKET_ERROR_CODE:
                    self.total_bytes = self.bytes_read
                    break
                else:
                    self.total_bytes = int(self.control_file.readline())
            if self.data_socket in readable:
                to_read = 1024 if length <= 0 else length - result.tell()
                if self.total_bytes and self.total_bytes < to_read:
                    to_read = self.total_bytes
                try:
                    data = self.data_socket.recv(to_read)
                    self.bytes_read += len(data)
                    result.write(data)
                except socket.error, e:
                    if e.errno == errno.EAGAIN:
                        pass
                    else:
                        raise
                if 0 < length == result.tell():
                    break
        if self.bytes_read == self.total_bytes:
            self.done = True
        return result.getvalue()

    def close(self):
        pass

class PipeControlledOutputSocket(PipeControlledBaseSocket):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        super(PipeControlledOutputSocket, self).__init__(data_file,
                                                         data_socket,
                                                         control_file,
                                                         control_socket)
        self.bytes_written = 0

    def write(self, data):
        view = memoryview(data)
        total_bytes_to_write = len(view)
        total_sent = 0
        while True:
            if total_sent == total_bytes_to_write:
                break
            _, writable, _ = select.select([], [self.data_socket], [], 0.5)
            if self.data_socket in writable:
                sent = self.data_socket.send(view[total_sent:])
                total_sent += sent
        self.bytes_written += total_bytes_to_write

    def close(self):
        self.control_socket.setblocking(1)
        length_control_message = '{}\r\n{}\r\n'.format(
            CONTROL_SOCKET_LENGTH_CODE, self.bytes_written)
        self.control_file.write(length_control_message)
        self.control_file.flush()

def do_regular_client_request(req_type, payload):
    return _do_client_request(req_type, payload)

def do_piped_client_request(req_type, payload):
    def pipe_handler(sock, req, res, uid):
        sock.setblocking(0)
        try:
            pipe_control_handler = PipeControlSocketHandler(uid)
            pipe_control_handler.connect()
            piped_output = PipeControlledOutputSocket(
                req,
                sock,
                pipe_control_handler.wfile,
                pipe_control_handler.conn)
            if _has_read_data(sys.stdin):
                shutil.copyfileobj(sys.stdin, piped_output)
            try:
                piped_output.close()
            except socket.error, e:
                if e.errno in [errno.EPIPE]:
                    raise ExecutionError('while closing piped output socket: '
                                         '{}'.format(e))
                else:
                    raise
            piped_input = PipeControlledInputSocket(
                res,
                sock,
                pipe_control_handler.rfile,
                pipe_control_handler.conn)
            shutil.copyfileobj(piped_input, sys.stdout)
            piped_input.close()
            pipe_control_handler.conn.setblocking(1)
            return pipe_control_handler
        finally:
            sock.setblocking(1)
    return _do_client_request(req_type, payload, pipe_handler)

def _do_client_request(req_type, payload, pipe_handler=None):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(unix_socket_path)
    except socket.error, e:
        if e.errno in [errno.ENOENT, errno.ECONNREFUSED]:
            raise ExecutionError('Is the pysource daemon running?'
                                 ' Run "pysource daemon start" to start it.')
        else:
            raise

    req = sock.makefile('wb', 0)
    res = sock.makefile('rb', -1)
    piped = pipe_handler is not None
    uid = str(uuid.uuid4())
    try:
        _write_body(req, {
            'req_type': req_type,
            'payload': payload,
            'piped': piped,
            'uid': uid
        })
        if piped:
            pipe_control_handler = pipe_handler(sock, req, res, uid)
            res_body = _read_body(pipe_control_handler.rfile)
            pipe_control_handler.close()
        else:
            res_body = _read_body(res)
        res_body_payload = res_body['payload']
        if res_body['status'] == RESPONSE_STATUS_ERROR:
            error = res_body_payload['error']
            raise ExecutionError('{0}'.format(error))
        return res_body_payload
    finally:
        req.close()
        res.close()
        sock.close()

def _read_body(sock):
    json_body_len = int(sock.readline())
    return json.loads(sock.read(json_body_len))

def _write_body(sock, body):
    json_body = json.dumps(body)
    json_body_len = len(json_body)
    sock.write(json_body_len)
    sock.write('\r\n')
    sock.write(json_body)

def _has_read_data(sock):
    return sock in select.select([sock], [], [], 0.0)[0]

def start_server():
    server = ThreadingUnixStreamServer(unix_socket_path,
                                       RequestHandler)
    server.serve_forever()

def cleanup():
    """Used for forced cleanup"""
    if os.path.exists(unix_socket_path):
        os.remove(unix_socket_path)
<commit_msg>continue work on making pipes work. not going well so far<commit_after># Copyright 2014 Dan Kilman
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import time
import uuid
import select
import os
import json
import socket
import errno
from SocketServer import (ThreadingUnixStreamServer,
                          StreamRequestHandler)
from io import BytesIO

from pysource import env
from pysource import ExecutionError
from pysource import remote_call_handlers
from pysource import request_context

RESPONSE_STATUS_OK = "ok"
RESPONSE_STATUS_ERROR = "error"

CONTROL_SOCKET_ERROR_CODE = 1
CONTROL_SOCKET_LENGTH_CODE = 2

unix_socket_path = os.path.join(env.pysource_dir, 'socket')

DEBUG = False
DEBUG_CLIENT = True
DEBUG_SERVER = False

def _handle(req_type, payload, **kwargs):
    handler = remote_call_handlers.get(req_type)
    if handler is None:
        raise RuntimeError('Unknown request type: {0}'.format(req_type))
    return handler(**payload)

class RequestHandler(StreamRequestHandler):

    def __init__(self, request, client_address, server):
        self.pipe_control_handler = None
        StreamRequestHandler.__init__(self, request, client_address, server)

    def handle(self):
        try:
            res_status = RESPONSE_STATUS_OK
            body = _read_body(self.rfile)
            if body['piped'] is True:
                request_context.piped = True
                self.pipe_control_handler = \
                    PipeControlSocketHandler(body['uid'])
                self._handle_piped(body, self.pipe_control_handler)
                res_payload = {}
            else:
                res_payload = _handle(**body)
        except Exception, e:
            res_status = RESPONSE_STATUS_ERROR
            res_payload = {'error': str(e)}

        if self.pipe_control_handler:
            response_file = self.pipe_control_handler.wfile
            if res_status == RESPONSE_STATUS_ERROR:
                try:
                    self.wfile.flush()
                except socket.error:
                    pass
                response_file.write('{}\r\n'.format(CONTROL_SOCKET_ERROR_CODE))
                response_file.flush()
        else:
            response_file = self.wfile
        _write_body(response_file, {
            'payload': res_payload,
            'status': res_status
        })

    def finish(self):
        StreamRequestHandler.finish(self)
        if self.pipe_control_handler:
            self.pipe_control_handler.close()

    def _handle_piped(self, body, pipe_control_handler):
        self.connection.setblocking(0)
        try:
            pipe_control_handler.accept()
            piped_input = PipeControlledInputSocket(
                self.rfile,
                self.connection,
                pipe_control_handler.rfile,
                pipe_control_handler.conn)
            piped_output = PipeControlledOutputSocket(
                self.wfile,
                self.connection,
                pipe_control_handler.wfile,
                pipe_control_handler.conn)
            request_context.stdin = piped_input
            request_context.stdout = piped_output
            result = _handle(**body)
            if result:
                piped_output.write(str(result))
            piped_input.close()
            piped_output.close()
        finally:
            self.connection.setblocking(1)
            pipe_control_handler.conn.setblocking(1)

class PipeControlSocketHandler(object):

    def __init__(self, uid):
        self.uid = uid
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.unix_socket_path = os.path.join(env.pysource_dir, self.uid)
        self.conn = None
        self.wfile = None
        self.rfile = None
        self.server = True

    def accept(self):
        self.sock.bind(self.unix_socket_path)
        self.sock.listen(1)
        conn, addr = self.sock.accept()
        self._setup(conn)

    def connect(self):
        self.server = False
        retries = 3
        current_tries = 0
        while True:
            try:
                current_tries += 1
                self.sock.connect(self.unix_socket_path)
                break
            except socket.error, e:
                if e.errno == errno.ENOENT and current_tries < retries:
                    time.sleep(0.1)
                else:
                    raise

        self._setup(self.sock)

    def _setup(self, conn):
        self.conn = conn
        self.conn.setblocking(0)
        self.rfile = self.conn.makefile('rb', -1)
        self.wfile = self.conn.makefile('wb', 0)

    def close(self):
        if not self.wfile.closed:
            try:
                self.conn.setblocking(1)
                self.wfile.flush()
            except socket.error:
                pass
        self.wfile.close()
        self.rfile.close()
        if self.server:
            try:
                self.conn.shutdown(socket.SHUT_WR)
            except socket.error:
                pass
            try:
                self.sock.close()
            except socket.error:
                pass
        self.conn.close()
        if self.server and os.path.exists(self.unix_socket_path):
            os.remove(self.unix_socket_path)

class PipeControlledBaseSocket(object):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        self.closed = False
        self.data_file = data_file
        self.data_socket = data_socket
        self.control_file = control_file
        self.control_socket = control_socket
        self.sockets = [self.data_socket, self.control_socket]

class PipeControlledInputSocket(PipeControlledBaseSocket):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        super(PipeControlledInputSocket, self).__init__(data_file,
                                                        data_socket,
                                                        control_file,
                                                        control_socket)
        self.bytes_read = 0
        self.total_bytes = None
        self.done = False

    def p(self, message):
        if DEBUG and ((DEBUG_CLIENT and not request_context.piped) or (DEBUG_SERVER and request_context.piped)):
            prefix = '(SERVER_IN)' if request_context.piped else '(CLIENT_IN)'
            print '{} {} (total_bytes={}, bytes_read={}, done={})'.format(prefix, message, self.total_bytes, self.bytes_read, self.done)

    def read(self, length=0, ensure_length=True):
        self.p('read length={}, ensure_length={}'.format(length, ensure_length))
        if self.done:
            return None
        result = BytesIO()
        while True:
            self.p('loop')
            if self.bytes_read == self.total_bytes:
                break
            readable, _, _ = select.select(self.sockets, [], [], 0.5)
            if self.control_socket in readable:
                control_code = int(self.control_file.readline())
                self.p('control_code={}'.format(control_code))
                if control_code == CONTROL_SOCKET_ERROR_CODE:
                    self.total_bytes = self.bytes_read
                    break
                else:
                    self.total_bytes = int(self.control_file.readline())
                self.p('total_bytes={}'.format(self.total_bytes))
            if self.data_socket in readable:
                to_read = 1024 if length <= 0 else length - result.tell()
                if self.total_bytes and self.total_bytes < to_read:
                    to_read = self.total_bytes
                try:
                    data = self.data_socket.recv(to_read)
                    self.p('data={}, len={}'.format(data, len(data)))
                    self.bytes_read += len(data)
                    result.write(data)
                except socket.error, e:
                    if e.errno == errno.EAGAIN:
                        self.p('errno.EAGAIN ensure_length={}'.format(ensure_length))
                        if ensure_length:
                            pass
                        else:
                            break
                    else:
                        raise
                if 0 < length == result.tell():
                    break
            if len(readable) == 0 and not ensure_length:
                break
        if self.bytes_read == self.total_bytes:
            self.done = True
        self.p('out done={}, bytes_read={}, total_bytes={}'.format(self.done, self.bytes_read, self.total_bytes))
        return result.getvalue()

    def close(self):
        self.p('read_close done={}, bytes_read={}, total_bytes={}'.format(self.done, self.bytes_read, self.total_bytes))
        pass

class PipeControlledOutputSocket(PipeControlledBaseSocket):

    def __init__(self, data_file, data_socket, control_file, control_socket):
        super(PipeControlledOutputSocket, self).__init__(data_file,
                                                         data_socket,
                                                         control_file,
                                                         control_socket)
        self.bytes_written = 0

    def p(self, message):
        if DEBUG and ((DEBUG_CLIENT and not request_context.piped) or (DEBUG_SERVER and request_context.piped)):
            prefix = '(SERVER_OUT)' if request_context.piped else '(CLIENT_OUT)'
            print '{} {}'.format(prefix, message)

    def write(self, data):
        self.p('write data={}'.format(data))
        view = memoryview(data)
        total_bytes_to_write = len(view)
        total_sent = 0
        while True:
            if total_sent == total_bytes_to_write:
                break
            _, writable, _ = select.select([], [self.data_socket], [], 0.5)
            if self.data_socket in writable:
                sent = self.data_socket.send(view[total_sent:])
                total_sent += sent
        self.bytes_written += total_bytes_to_write
        self.p('write out bytes_written={}'.format(self.bytes_written))

    def close(self):
        self.p('write close bytes_written={}'.format(self.bytes_written))
        self.control_socket.setblocking(1)
        try:
            length_control_message = '{}\r\n{}\r\n'.format(
                CONTROL_SOCKET_LENGTH_CODE, self.bytes_written)
            self.control_file.write(length_control_message)
            self.control_file.flush()
        finally:
            self.control_socket.setblocking(0)

def do_regular_client_request(req_type, payload):
    return _do_client_request(req_type, payload)

def do_piped_client_request(req_type, payload):

    def p(message):
        if DEBUG:
            print '(CLIENT) {}'.format(message)

    def pipe_handler(sock, req, res, uid):
        sock.setblocking(0)
        try:
            pipe_control_handler = PipeControlSocketHandler(uid)
            pipe_control_handler.connect()
            piped_output = PipeControlledOutputSocket(
                req,
                sock,
                pipe_control_handler.wfile,
                pipe_control_handler.conn)
            piped_input = PipeControlledInputSocket(
                res,
                sock,
                pipe_control_handler.rfile,
                pipe_control_handler.conn)

            stdin_buf = 16 * 1024
            data_sock_buf = 1024
            stdin = sys.stdin
            data_sock = sock
            control_soc = pipe_control_handler.conn
            read_sockets = [stdin, data_sock, control_soc]
            control_soc_written = False
            stdin_done = False
            stdin_done_handled = False

            count = 0
            while True:
                if count % 1000 == 0:
                    p('loop count={}'.format(count))
                count += 1
                readable = select.select(read_sockets, [], [], 0.01)[0]
                if stdin in readable:
                    buf = stdin.read(stdin_buf)
                    if not buf:
                        stdin_done = True
                    else:
                        p('stdin readable')
                        piped_output.write(buf)
                        p('write_buf ended')
                else:
                    stdin_done = True
                if not stdin_done_handled and stdin_done:
                    piped_output.close()
                    p('piped_output_close')
                    read_sockets = [control_soc, data_sock]
                    stdin_done_handled = True
                if control_soc in readable:
                    control_soc_written = True
                    p('control_sock_written=True')
                if data_sock in readable or control_soc in readable:
                    if control_soc_written:
                        p('piped_input_read()')
                        buf = piped_input.read()
                    else:
                        p('piped_input_read(data_sock_buf)')
                        buf = piped_input.read(data_sock_buf,
                                               ensure_length=False)
                    if buf is None:
                        p('buf is None')
                        break
                    p('read ended len={}'.format(len(buf)))
                    if len(buf) > 0:
                        p('stdout.write len={}'.format(len(buf)))
                        sys.stdout.write(buf)
                        try:
                            sys.stdout.flush()
                        except IOError, e:
                            if e.errno == errno.EPIPE:
                                raise ExecutionError('Flushing stdout failed.'
                                                     ' It seems the process'
                                                     ' being piped to, '
                                                     'terminated.')
                            else:
                                raise
                    else:
                        p('stdout.write NO_WRITE')
                if control_soc_written:
                    break

            p('piped_input_close')
            piped_input.close()
            pipe_control_handler.conn.setblocking(1)
            return pipe_control_handler
        finally:
            sock.setblocking(1)
    return _do_client_request(req_type, payload, pipe_handler)

def _do_client_request(req_type, payload, pipe_handler=None):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(unix_socket_path)
    except socket.error, e:
        if e.errno in [errno.ENOENT, errno.ECONNREFUSED]:
            raise ExecutionError('Is the pysource daemon running?'
                                 ' Run "pysource daemon start" to start it.')
        else:
            raise

    req = sock.makefile('wb', 0)
    res = sock.makefile('rb', -1)
    piped = pipe_handler is not None
    uid = str(uuid.uuid4())
    try:
        _write_body(req, {
            'req_type': req_type,
            'payload': payload,
            'piped': piped,
            'uid': uid
        })
        if piped:
            pipe_control_handler = pipe_handler(sock, req, res, uid)
            res_body = _read_body(pipe_control_handler.rfile)
            pipe_control_handler.close()
        else:
            res_body = _read_body(res)
        res_body_payload = res_body['payload']
        if res_body['status'] == RESPONSE_STATUS_ERROR:
            error = res_body_payload['error']
            raise ExecutionError('{0}'.format(error))
        return res_body_payload
    finally:
        req.close()
        res.close()
        sock.close()

def _read_body(sock):
    json_body_len = int(sock.readline())
    return json.loads(sock.read(json_body_len))

def _write_body(sock, body):
    json_body = json.dumps(body)
    json_body_len = len(json_body)
    sock.write(json_body_len)
    sock.write('\r\n')
    sock.write(json_body)

def start_server():
    server = ThreadingUnixStreamServer(unix_socket_path,
                                       RequestHandler)
    server.serve_forever()

def cleanup():
    """Used for forced cleanup"""
    if os.path.exists(unix_socket_path):
        os.remove(unix_socket_path)
<|endoftext|>


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

import sys
import re
import shlex

import genutil

def _die(s):
    genutil.die(s)
def _msgb(b,s=''):
    genutil.msgb(b,s)

class map_info_t(object):
    def __init__(self):
        self.map_name = None
        self.space = None  # legacy, vex, evex, xop
        self.legacy_escape = None # N/A or 0f
        self.legacy_opcode = None # N/A or 38, 3a
        self.map_id = None  # N/A or 0,1,2,3,... 8,9,0xA
        # "var" means variable, requires a table generated based on defined instructions
        self.modrm = None  # var,yes,no, has modrm
        self.imm8 = None   # var,yes,no, has imm8
        self.imm32 = None  # var,yes,no, has imm32
        self.opcpos = None  # 0,1,2, ... -1 (last) opcode position in pattern
        self.priority = 10
        # search_pattern is the string that we use to identify this
        # map in the XED decode patterns. The pattern may have spaces
        # in it. (and motivates using shlex to parse the input lines)
        self.search_pattern = None

    def is_legacy(self):
        return self.space == 'legacy'
    def is_vex(self):
        return self.space == 'vex'
    def is_evex(self):
        return self.space == 'evex'
    def is_xop(self):
        return self.space == 'xop'

    def map_short_name(self):
        if self.map_name  == 'amd-3dnow':
            return 'AMD'
        h = hex(self.map_id)[-1]
        return str(h)
    def ild_enum(self):
        s = self.map_short_name()
        if self.space == 'XOP':
            s = '_XOP{}'.format(s)
        return 'XED_ILD_MAP{}'.format(s)

    def get_legacy_escapes(self):
        if self.legacy_opcode == 'N/A':
            return (self.legacy_escape, None)
        return (self.legacy_escape, self.legacy_opcode)

    def has_variable_modrm(self):
        return self.modrm == 'var'
    def has_regular_modrm(self):
        return self.modrm == 'yes'

    def has_variable_imm8(self):
        return self.imm8 == 'var'
    def has_regular_imm8(self):
        return self.imm8 == 'yes'

    def has_variable_imm32(self):
        return self.imm32 == 'var'
    def has_regular_imm32(self):
        return self.imm32 == 'yes'

    def __str__(self):
        s = []
        s.append("name: {}".format(self.map_name))
        s.append("space: {}".format(self.space))
        s.append("legacyesc: {}".format(self.legacy_escape))
        s.append("legacyopc: {}".format(self.legacy_opcode))
        s.append("mapid: {}".format(self.map_id))
        s.append("modrm: {}".format(self.modrm))
        s.append("imm8: {}".format(self.imm8))
        s.append("imm32: {}".format(self.imm32))
        s.append("opcpos: {}".format(self.opcpos))
        s.append("priority: {}".format(self.priority))
        s.append("search_pattern: {}".format(self.search_pattern))
        return " ".join(s)

def _parse_map_line(s):
    # shlex allows for quoted substrings containing spaces as
    # individual args.
    t = shlex.split(s.strip())
    if len(t) != 10:
        _die("Bad map description line: [{}]".format(s))
    mi = map_info_t()
    mi.map_name = t[0]
    mi.space = t[1]
    mi.legacy_escape = t[2]
    mi.legacy_opcode = t[3]
    mi.map_id = t[4]
    mi.modrm = t[5]
    mi.imm8 = t[6]
    mi.imm32 = t[7]
    mi.opcpos = t[8]
    mi.search_pattern = t[9]

    if mi.space not in ['legacy','vex','evex', 'xop']:
        _die("Bad map description encoding space [{}]".format(s))
    if mi.space == 'legacy':
        if genutil.is_hex(mi.legacy_escape):
            pass
        elif mi.legacy_escape != 'N/A':
            _die("Bad map description legacy escape [{}]".format(s))
        if  genutil.is_hex(mi.legacy_opcode):
            pass
        elif mi.legacy_opcode != 'N/A':
            _die("Bad map description legacy opcode [{}]".format(s))

        if mi.map_id == 'N/A':
            _die("Bad map description map-id [{}]".format(s))
    else:
        if mi.legacy_escape != 'N/A':
            _die("Bad map description legacy escape [{}]".format(s))
        if mi.legacy_opcode != 'N/A':
            _die("Bad map description legacy opcode [{}]".format(s))
        if genutil.numeric(mi.map_id):
            mi.map_id = genutil.make_numeric(mi.map_id)
        else:
            _die("Bad map description map id [{}]".format(s))

    if mi.modrm not in ['var','yes','no']:
        _die("Bad map description modrm specifier [{}]".format(s))
    if mi.imm8 not in ['var','yes','no']:
        _die("Bad map description imm8 specifier [{}]".format(s))
    if mi.imm32 not in ['var','yes','no']:
        _die("Bad map description imm32 specifier [{}]".format(s))
    if genutil.numeric(mi.opcpos):
        mi.opcpos = genutil.make_numeric(mi.opcpos)
    else:
        _die("Bad map description opcode position specifier [{}]".format(s))

    # we want the longer patterns first when we sort the map_info_t.
    mi.priority = 100-len(mi.search_pattern)
    return mi

def read_file(fn):
    lines = open(fn,'r').readlines()
    lines = map(genutil.no_comments, lines)
    lines = list(filter(genutil.blank_line, lines))
    maps = [] # list of map_info_t
    for line in lines:
        maps.append( _parse_map_line(line) )
    maps.sort(key=lambda x: x.priority)
    for m in maps:
        _msgb("MAPINFO",m)
    return maps

if __name__ == "__main__":
    read_file(sys.argv[1])
    sys.exit(0)

<commit_msg>map_info_rdr.py: use stderr<commit_after>#!/usr/bin/env python
# -*- python -*-
#BEGIN_LEGAL
#
#Copyright (c) 2019 Intel Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#END_LEGAL

import sys
import re
import shlex

import genutil

def _die(s):
    genutil.die(s)
def _msgb(b,s=''):
    genutil.msgerr("[{}] {}".format(b,s))

class map_info_t(object):
    def __init__(self):
        self.map_name = None
        self.space = None  # legacy, vex, evex, xop
        self.legacy_escape = None # N/A or 0f
        self.legacy_opcode = None # N/A or 38, 3a
        self.map_id = None  # N/A or 0,1,2,3,... 8,9,0xA
        # "var" means variable, requires a table generated based on defined instructions
        self.modrm = None  # var,yes,no, has modrm
        self.imm8 = None   # var,yes,no, has imm8
        self.imm32 = None  # var,yes,no, has imm32
        self.opcpos = None  # 0,1,2, ... -1 (last) opcode position in pattern
        self.priority = 10
        # search_pattern is the string that we use to identify this
        # map in the XED decode patterns. The pattern may have spaces
        # in it. (and motivates using shlex to parse the input lines)
        self.search_pattern = None

    def is_legacy(self):
        return self.space == 'legacy'
    def is_vex(self):
        return self.space == 'vex'
    def is_evex(self):
        return self.space == 'evex'
    def is_xop(self):
        return self.space == 'xop'

    def map_short_name(self):
        if self.map_name  == 'amd-3dnow':
            return 'AMD'
        h = hex(self.map_id)[-1]
        return str(h)
    def ild_enum(self):
        s = self.map_short_name()
        if self.space == 'XOP':
            s = '_XOP{}'.format(s)
        return 'XED_ILD_MAP{}'.format(s)

    def get_legacy_escapes(self):
        if self.legacy_opcode == 'N/A':
            return (self.legacy_escape, None)
        return (self.legacy_escape, self.legacy_opcode)

    def has_variable_modrm(self):
        return self.modrm == 'var'
    def has_regular_modrm(self):
        return self.modrm == 'yes'

    def has_variable_imm8(self):
        return self.imm8 == 'var'
    def has_regular_imm8(self):
        return self.imm8 == 'yes'

    def has_variable_imm32(self):
        return self.imm32 == 'var'
    def has_regular_imm32(self):
        return self.imm32 == 'yes'

    def __str__(self):
        s = []
        s.append("name: {}".format(self.map_name))
        s.append("space: {}".format(self.space))
        s.append("legacyesc: {}".format(self.legacy_escape))
        s.append("legacyopc: {}".format(self.legacy_opcode))
        s.append("mapid: {}".format(self.map_id))
        s.append("modrm: {}".format(self.modrm))
        s.append("imm8: {}".format(self.imm8))
        s.append("imm32: {}".format(self.imm32))
        s.append("opcpos: {}".format(self.opcpos))
        s.append("priority: {}".format(self.priority))
        s.append("search_pattern: {}".format(self.search_pattern))
        return " ".join(s)

def _parse_map_line(s):
    # shlex allows for quoted substrings containing spaces as
    # individual args.
    t = shlex.split(s.strip())
    if len(t) != 10:
        _die("Bad map description line: [{}]".format(s))
    mi = map_info_t()
    mi.map_name = t[0]
    mi.space = t[1]
    mi.legacy_escape = t[2]
    mi.legacy_opcode = t[3]
    mi.map_id = t[4]
    mi.modrm = t[5]
    mi.imm8 = t[6]
    mi.imm32 = t[7]
    mi.opcpos = t[8]
    mi.search_pattern = t[9]

    if mi.space not in ['legacy','vex','evex', 'xop']:
        _die("Bad map description encoding space [{}]".format(s))
    if mi.space == 'legacy':
        if genutil.is_hex(mi.legacy_escape):
            pass
        elif mi.legacy_escape != 'N/A':
            _die("Bad map description legacy escape [{}]".format(s))
        if  genutil.is_hex(mi.legacy_opcode):
            pass
        elif mi.legacy_opcode != 'N/A':
            _die("Bad map description legacy opcode [{}]".format(s))

        if mi.map_id == 'N/A':
            _die("Bad map description map-id [{}]".format(s))
    else:
        if mi.legacy_escape != 'N/A':
            _die("Bad map description legacy escape [{}]".format(s))
        if mi.legacy_opcode != 'N/A':
            _die("Bad map description legacy opcode [{}]".format(s))
        if genutil.numeric(mi.map_id):
            mi.map_id = genutil.make_numeric(mi.map_id)
        else:
            _die("Bad map description map id [{}]".format(s))

    if mi.modrm not in ['var','yes','no']:
        _die("Bad map description modrm specifier [{}]".format(s))
    if mi.imm8 not in ['var','yes','no']:
        _die("Bad map description imm8 specifier [{}]".format(s))
    if mi.imm32 not in ['var','yes','no']:
        _die("Bad map description imm32 specifier [{}]".format(s))
    if genutil.numeric(mi.opcpos):
        mi.opcpos = genutil.make_numeric(mi.opcpos)
    else:
        _die("Bad map description opcode position specifier [{}]".format(s))

    # we want the longer patterns first when we sort the map_info_t.
    mi.priority = 100-len(mi.search_pattern)
    return mi

def read_file(fn):
    lines = open(fn,'r').readlines()
    lines = map(genutil.no_comments, lines)
    lines = list(filter(genutil.blank_line, lines))
    maps = [] # list of map_info_t
    for line in lines:
        maps.append( _parse_map_line(line) )
    maps.sort(key=lambda x: x.priority)
    for m in maps:
        _msgb("MAPINFO",m)
    return maps

if __name__ == "__main__":
    read_file(sys.argv[1])
    sys.exit(0)

<|endoftext|>


<commit_before>"""Methods for reading files.
"""

import numpy as np

def read_timeseries(inputfns, importer, **kwargs):
    """Read a list of input files using io tools and stack them into a 3d array.

    Parameters
    ----------
    inputfns : list
        List of input files returned by any function implemented in archive.
    importer : function
        Any function implemented in importers.
    kwargs : dict
        Optional keyword arguments for the importer.

    Returns
    -------
    out : tuple
        A three-element tuple containing the precipitation fields read, the quality fields,
        and associated metadata.

    """

    # check for missing data
    Rref = None
    if all(ifn is None for ifn in inputfns):
        return None, None, None
    else:
        if len(inputfns[0]) == 0:
            return None, None, None
        for ifn in inputfns[0]:
            if ifn is not None:
                Rref, Qref, metadata = importer(ifn, **kwargs)
                break

    if Rref is None:
        return None, None, None

    R = []
    Q = []
    timestamps = []
    for i,ifn in enumerate(inputfns[0]):
        if ifn is not None:
            R_, Q_, _ = importer(ifn, **kwargs)
            R.append(R_)
            Q.append(Q_)
            timestamps.append(inputfns[1][i])
        else:
            R.append(Rref*np.nan)
            if Qref is not None:
                Q.append(Qref*np.nan)
            else:
                Q.append(None)
            timestamps.append(inputfns[1][i])

    # Replace this with stack?
    R = np.concatenate([R_[None, :, :] for R_ in R])
    #TODO: Q should be organized as R, but this is not trivial as Q_ can be also None or a scalar
    metadata["timestamps"] = np.array(timestamps)

    return R, Q, metadata
<commit_msg>Docstring revisions<commit_after>"""Methods for reading files.
"""

import numpy as np

def read_timeseries(inputfns, importer, **kwargs):
    """Read a time series of input files using the methods implemented in the
    importers module and stack them into a 3d array of shape (num_timesteps,
    height, width).

    Parameters
    ----------
    inputfns : tuple
        Input files returned by a function implemented in the archive module.
    importer : function
        A function implemented in the importers module.
    kwargs : dict
        Optional keyword arguments for the importer.

    Returns
    -------
    out : tuple
        A three-element tuple containing the read data and quality rasters and
        associated metadata. If an input file name is None, the corresponding
        precipitation and quality fields are filled with nan values. If all
        input file names are None or if the length of the file name list is
        zero, a three-element tuple containing None values is returned.

    """

    # check for missing data
    Rref = None
    if all(ifn is None for ifn in inputfns):
        return None, None, None
    else:
        if len(inputfns[0]) == 0:
            return None, None, None
        for ifn in inputfns[0]:
            if ifn is not None:
                Rref, Qref, metadata = importer(ifn, **kwargs)
                break

    if Rref is None:
        return None, None, None

    R = []
    Q = []
    timestamps = []
    for i,ifn in enumerate(inputfns[0]):
        if ifn is not None:
            R_, Q_, _ = importer(ifn, **kwargs)
            R.append(R_)
            Q.append(Q_)
            timestamps.append(inputfns[1][i])
        else:
            R.append(Rref*np.nan)
            if Qref is not None:
                Q.append(Qref*np.nan)
            else:
                Q.append(None)
            timestamps.append(inputfns[1][i])

    # Replace this with stack?
    R = np.concatenate([R_[None, :, :] for R_ in R])
    #TODO: Q should be organized as R, but this is not trivial as Q_ can be also None or a scalar
    metadata["timestamps"] = np.array(timestamps)

    return R, Q, metadata
<|endoftext|>


<commit_before>#!/usr/bin/env python
import sys
import cmd
from docopt import docopt, DocoptExit
from pytelemetry import Pytelemetry
import pytelemetry.transports.serialtransport as transports
from pytelemetrycli.topics import Topics
from pytelemetrycli.runner import Runner
from serial.tools import list_ports
from pytelemetrycli.ui.superplot import Superplot, PlotType
from threading import Lock
from pytelemetrycli.initialization import init_logging
import logging
from logging import getLogger
import os

logger = getLogger('cli')

def docopt_cmd(func):
    def fn(self, arg):
        try:
            if fn.__name__ == "do_pub":
                # Fix for negative numbers
                opt = docopt(fn.__doc__, arg, options_first=True)
            else:
                opt = docopt(fn.__doc__, arg)

        except DocoptExit as e:
            print('Command is invalid. See :')
            print(e)
            return

        except SystemExit:
            # The SystemExit exception prints the usage for --help
            # We do not need to do the print here.
            return

        return func(self, opt)

    fn.__name__ = func.__name__
    fn.__doc__ = func.__doc__
    fn.__dict__.update(func.__dict__)
    return fn

class Application (cmd.Cmd):
    def __init__(self, transport=None, stdout=None):
        # cmd Initialization and configuration
        cmd.Cmd.__init__(self,stdout=stdout)
        self.intro = 'pytelemetry terminal started.' \
                 + ' (type help for a list of commands.)'
        self.prompt = ':> '
        self.file = None

        # pytelemetry setup
        if not transport:
            self.transport = transports.SerialTransport()
        else:
            self.transport = transport
        self.telemetry = Pytelemetry(self.transport)
        self.topics = Topics()
        self.plots = []
        self.plotsLock = Lock()
        self.runner = Runner(self.transport,
                             self.telemetry,
                             self.plots,
                             self.plotsLock,
                             self.topics)
        self.telemetry.subscribe(None,self.topics.process)

        self.types_lookup = {'--s'    :  'string',
                             '--u8'   :  'uint8',
                             '--u16'  :  'uint16',
                             '--u32'  :  'uint32',
                             '--i8'   :  'int8',
                             '--i16'  :  'int16',
                             '--i32'  :  'int32',
                             '--f32'  :  'float32'}
        logger.info("Module path : %s" % os.path.dirname(os.path.realpath(__file__)))
        try:
            logger.info("Module version : %s" % __version__)
        except:
            logger.warning("Module version : not found.")

    @docopt_cmd
    def do_serial(self, arg):
        """
Connects pytelemetry to the serial port.

Usage: serial <port> [options]

Options:
-b X, --bauds X        Connection speed in bauds  [default: 9600]

        """
        try:
            self.runner.disconnect()
            logger.warn("User requested connect without desconnecting first.")
        except (IOError,AttributeError) as e:
            logger.warn("Already disconnected. Continuing happily. E : {0}"
                         .format(e))
            pass

        try:
            b = int(arg['--bauds'])
            self.runner.connect(arg['<port>'],b)
        except IOError as e:
            self.stdout.write("Failed to connect to {0} at {1} (bauds).\n"
                    .format(arg['<port>'],b))

            logger.warn("Failed to connect to {0} at {1} (bauds). E : \n"
                          .format(arg['<port>'],b,e))
        else:
            s = "Connected to {0} at {1} (bauds).\n".format(arg['<port>'],b)
            self.stdout.write(s)
            logger.info(s)
            self.topics.clear()
            logger.info("Cleared all topics for new session.")

    @docopt_cmd
    def do_print(self, arg):
        """
Prints X last received samples from <topic>.

Usage: print <topic> [options]

Options:
-a X, --amount X        Amount of samples to display [default: 1]

        """
        topic = arg['<topic>']
        if not self.topics.exists(topic):
            s = "Topic '{0}' unknown. Type 'ls' to list all available topics.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        try:
            amount = int(arg['--amount'])
        except:
            s = "Could not cast --amount = '{0}' to integer. Using 1.\n".format(arg['--amount'])
            self.stdout.write(s)
            logger.warn(s)
            amount = 1

        s = self.topics.samples(topic,amount)

        if s is not None:
            for i in s:
                self.stdout.write("%i\n" % i)
        else:
            logger.error("Could not retrieve {0} sample(s) under topic '{1}'.\n".format(amount,topic))

    @docopt_cmd
    def do_ls(self, arg):
        """
Without options, prints a list of all received topics.
With the --serial flag, prints a list of all available COM ports

Usage: ls [options]

Options:
-s, --serial     Use this flag to print a list of all available serial ports

        """
        if arg['--serial']:
            self.stdout.write("Available COM ports:\n")
            for port,desc,hid in list_ports.comports():
                self.stdout.write("%s \t %s\n" % (port,desc))
        else:

            for topic in self.topics.ls():
                self.stdout.write("%s\n" % topic)

    @docopt_cmd
    def do_plot(self, arg):
        """
Plots <topic> in a graph window.

Usage: plot <topic>
        """

        topic = arg['<topic>']

        if not self.topics.exists(topic):
            s = "Topic '{0}' unknown. Type 'ls' to list all available topics.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        if self.topics.intransfer(topic):
            s = "Topic '{0}' already plotting.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        plotTypeFlag = self.topics.xytype(arg['<topic>'])
        plotType = PlotType.linear

        if plotTypeFlag == 'indexed':
            plotType = PlotType.indexed

        p = Superplot(topic,plotType)
        q, ctrl = p.start()

        # Protect self.plots from modifications from the runner thread
        self.plotsLock.acquire()

        self.plots.append({
            'topic': topic,
            'plot': p,     # Plot handler
            'queue': q,    # Data queue
            'ctrl': ctrl   # Plot control pipe
        })

        self.plotsLock.release()

        self.topics.transfer(topic,q)

        s = "Plotting '{0}' in mode [{1}].\n".format(topic,plotTypeFlag)
        logger.info(s)
        self.stdout.write(s)

    @docopt_cmd
    def do_pub(self, arg):
        """
Publishes a (value | string) on <topic>.

Usage: pub (--u8 | --u16 | --u32 | --i8 | --i16 | --i32 | --f32 | --s) <topic> <value>
        """

        if arg['--f32']:
            arg['<value>'] = float(arg['<value>'])
        elif not arg['--s']:
            arg['<value>'] = int(arg['<value>'])

        subset = {k: arg[k] for k in ("--u8","--u16","--u32","--i8","--i16","--i32","--f32","--s")}

        valtype = None
        for i, k in subset.items():
            if k:
                valtype = self.types_lookup[i]

        if not valtype:
            logger.error(
                "Payload type [{0}] unkown."
                .format(arg))
            return

        self.telemetry.publish(arg['<topic>'],arg['<value>'],valtype)

        s = "Published on topic '{0}' : {1} [{2}]\n".format(arg['<topic>'], arg['<value>'],valtype)
        self.stdout.write(s)
        logger.info(s)

    @docopt_cmd
    def do_count(self, arg):
        """
Prints a count of received samples for each topic.

Usage: count
        """
        for topic in self.topics.ls():
            self.stdout.write("{0} : {1}\n".format(topic, self.topics.count(topic)))

    @docopt_cmd
    def do_disconnect(self, arg):
        """
Disconnects from any open connection.

Usage: disconnect
        """
        try:
            self.runner.disconnect()
            self.stdout.write("Disconnected.\n")
            logger.info("Disconnected.")
        except:
            logger.warn("Already disconnected. Continuing happily.")

    @docopt_cmd
    def do_info(self, arg):
        """
Prints out cli.py full path, module version.

Usage: info
        """
        self.stdout.write("- CLI path : %s\n" % os.path.dirname(os.path.realpath(__file__)))
        try:
            self.stdout.write("- version : %s\n" % __version__)
        except:
            self.stdout.write("- version : not found.\n")

    @docopt_cmd
    def do_stats(self, arg):
        """
Displays different metrics about the active transport (ex : serial port).
This allows you to know if for instance corrupted frames are received, what fraction
of the maximum baudrate is being used, etc.

Usage: info
        """
        measures = self.telemetry.stats()
        for key,item in zip(measures.keys(),measures.items()):
            print("%s : %s" % (key,item))

    def do_quit(self, arg):
        """
Exits the terminal application.

Usage: quit
        """
        self.runner.terminate()
        self.stdout.write("Good Bye!\n")
        logger.info("Application quit.")
        exit()

# Main function to start from script or from entry point
def pytlm():
    init_logging()
    try:
        Application().cmdloop()
    except SystemExit:
        pass
    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    pytlm()
<commit_msg>Bug fix (raises exception with strings)<commit_after>#!/usr/bin/env python
import sys
import cmd
from docopt import docopt, DocoptExit
from pytelemetry import Pytelemetry
import pytelemetry.transports.serialtransport as transports
from pytelemetrycli.topics import Topics
from pytelemetrycli.runner import Runner
from serial.tools import list_ports
from pytelemetrycli.ui.superplot import Superplot, PlotType
from threading import Lock
from pytelemetrycli.initialization import init_logging
import logging
from logging import getLogger
import os

logger = getLogger('cli')

def docopt_cmd(func):
    def fn(self, arg):
        try:
            if fn.__name__ == "do_pub":
                # Fix for negative numbers
                opt = docopt(fn.__doc__, arg, options_first=True)
            else:
                opt = docopt(fn.__doc__, arg)

        except DocoptExit as e:
            print('Command is invalid. See :')
            print(e)
            return

        except SystemExit:
            # The SystemExit exception prints the usage for --help
            # We do not need to do the print here.
            return

        return func(self, opt)

    fn.__name__ = func.__name__
    fn.__doc__ = func.__doc__
    fn.__dict__.update(func.__dict__)
    return fn

class Application (cmd.Cmd):
    def __init__(self, transport=None, stdout=None):
        # cmd Initialization and configuration
        cmd.Cmd.__init__(self,stdout=stdout)
        self.intro = 'pytelemetry terminal started.' \
                 + ' (type help for a list of commands.)'
        self.prompt = ':> '
        self.file = None

        # pytelemetry setup
        if not transport:
            self.transport = transports.SerialTransport()
        else:
            self.transport = transport
        self.telemetry = Pytelemetry(self.transport)
        self.topics = Topics()
        self.plots = []
        self.plotsLock = Lock()
        self.runner = Runner(self.transport,
                             self.telemetry,
                             self.plots,
                             self.plotsLock,
                             self.topics)
        self.telemetry.subscribe(None,self.topics.process)

        self.types_lookup = {'--s'    :  'string',
                             '--u8'   :  'uint8',
                             '--u16'  :  'uint16',
                             '--u32'  :  'uint32',
                             '--i8'   :  'int8',
                             '--i16'  :  'int16',
                             '--i32'  :  'int32',
                             '--f32'  :  'float32'}
        logger.info("Module path : %s" % os.path.dirname(os.path.realpath(__file__)))
        try:
            logger.info("Module version : %s" % __version__)
        except:
            logger.warning("Module version : not found.")

    @docopt_cmd
    def do_serial(self, arg):
        """
Connects pytelemetry to the serial port.

Usage: serial <port> [options]

Options:
-b X, --bauds X        Connection speed in bauds  [default: 9600]

        """
        try:
            self.runner.disconnect()
            logger.warn("User requested connect without desconnecting first.")
        except (IOError,AttributeError) as e:
            logger.warn("Already disconnected. Continuing happily. E : {0}"
                         .format(e))
            pass

        try:
            b = int(arg['--bauds'])
            self.runner.connect(arg['<port>'],b)
        except IOError as e:
            self.stdout.write("Failed to connect to {0} at {1} (bauds).\n"
                    .format(arg['<port>'],b))

            logger.warn("Failed to connect to {0} at {1} (bauds). E : \n"
                          .format(arg['<port>'],b,e))
        else:
            s = "Connected to {0} at {1} (bauds).\n".format(arg['<port>'],b)
            self.stdout.write(s)
            logger.info(s)
            self.topics.clear()
            logger.info("Cleared all topics for new session.")

    @docopt_cmd
    def do_print(self, arg):
        """
Prints X last received samples from <topic>.

Usage: print <topic> [options]

Options:
-a X, --amount X        Amount of samples to display [default: 1]

        """
        topic = arg['<topic>']
        if not self.topics.exists(topic):
            s = "Topic '{0}' unknown. Type 'ls' to list all available topics.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        try:
            amount = int(arg['--amount'])
        except:
            s = "Could not cast --amount = '{0}' to integer. Using 1.\n".format(arg['--amount'])
            self.stdout.write(s)
            logger.warn(s)
            amount = 1

        s = self.topics.samples(topic,amount)

        if s is not None:
            for i in s:
                self.stdout.write("{0}\n".format(i))
        else:
            logger.error("Could not retrieve {0} sample(s) under topic '{1}'.\n".format(amount,topic))

    @docopt_cmd
    def do_ls(self, arg):
        """
Without options, prints a list of all received topics.
With the --serial flag, prints a list of all available COM ports

Usage: ls [options]

Options:
-s, --serial     Use this flag to print a list of all available serial ports

        """
        if arg['--serial']:
            self.stdout.write("Available COM ports:\n")
            for port,desc,hid in list_ports.comports():
                self.stdout.write("%s \t %s\n" % (port,desc))
        else:

            for topic in self.topics.ls():
                self.stdout.write("%s\n" % topic)

    @docopt_cmd
    def do_plot(self, arg):
        """
Plots <topic> in a graph window.

Usage: plot <topic>
        """

        topic = arg['<topic>']

        if not self.topics.exists(topic):
            s = "Topic '{0}' unknown. Type 'ls' to list all available topics.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        if self.topics.intransfer(topic):
            s = "Topic '{0}' already plotting.\n".format(topic)
            self.stdout.write(s)
            logger.warn(s)
            return

        plotTypeFlag = self.topics.xytype(arg['<topic>'])
        plotType = PlotType.linear

        if plotTypeFlag == 'indexed':
            plotType = PlotType.indexed

        p = Superplot(topic,plotType)
        q, ctrl = p.start()

        # Protect self.plots from modifications from the runner thread
        self.plotsLock.acquire()

        self.plots.append({
            'topic': topic,
            'plot': p,     # Plot handler
            'queue': q,    # Data queue
            'ctrl': ctrl   # Plot control pipe
        })

        self.plotsLock.release()

        self.topics.transfer(topic,q)

        s = "Plotting '{0}' in mode [{1}].\n".format(topic,plotTypeFlag)
        logger.info(s)
        self.stdout.write(s)

    @docopt_cmd
    def do_pub(self, arg):
        """
Publishes a (value | string) on <topic>.

Usage: pub (--u8 | --u16 | --u32 | --i8 | --i16 | --i32 | --f32 | --s) <topic> <value>
        """

        if arg['--f32']:
            arg['<value>'] = float(arg['<value>'])
        elif not arg['--s']:
            arg['<value>'] = int(arg['<value>'])

        subset = {k: arg[k] for k in ("--u8","--u16","--u32","--i8","--i16","--i32","--f32","--s")}

        valtype = None
        for i, k in subset.items():
            if k:
                valtype = self.types_lookup[i]

        if not valtype:
            logger.error(
                "Payload type [{0}] unkown."
                .format(arg))
            return

        self.telemetry.publish(arg['<topic>'],arg['<value>'],valtype)

        s = "Published on topic '{0}' : {1} [{2}]\n".format(arg['<topic>'], arg['<value>'],valtype)
        self.stdout.write(s)
        logger.info(s)

    @docopt_cmd
    def do_count(self, arg):
        """
Prints a count of received samples for each topic.

Usage: count
        """
        for topic in self.topics.ls():
            self.stdout.write("{0} : {1}\n".format(topic, self.topics.count(topic)))

    @docopt_cmd
    def do_disconnect(self, arg):
        """
Disconnects from any open connection.

Usage: disconnect
        """
        try:
            self.runner.disconnect()
            self.stdout.write("Disconnected.\n")
            logger.info("Disconnected.")
        except:
            logger.warn("Already disconnected. Continuing happily.")

    @docopt_cmd
    def do_info(self, arg):
        """
Prints out cli.py full path, module version.

Usage: info
        """
        self.stdout.write("- CLI path : %s\n" % os.path.dirname(os.path.realpath(__file__)))
        try:
            self.stdout.write("- version : %s\n" % __version__)
        except:
            self.stdout.write("- version : not found.\n")

    @docopt_cmd
    def do_stats(self, arg):
        """
Displays different metrics about the active transport (ex : serial port).
This allows you to know if for instance corrupted frames are received, what fraction
of the maximum baudrate is being used, etc.

Usage: info
        """
        measures = self.telemetry.stats()
        for key,item in zip(measures.keys(),measures.items()):
            print("%s : %s" % (key,item))

    def do_quit(self, arg):
        """
Exits the terminal application.

Usage: quit
        """
        self.runner.terminate()
        self.stdout.write("Good Bye!\n")
        logger.info("Application quit.")
        exit()

# Main function to start from script or from entry point
def pytlm():
    init_logging()
    try:
        Application().cmdloop()
    except SystemExit:
        pass
    except KeyboardInterrupt:
        pass

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