<commit_before>#! /usr/bin/python

import sys
import subprocess
import os
import time
from getIndoorTemp import getIndoorTemp
import RPi.GPIO as GPIO

DEBUG = 0

METRIC=0 # 0 for F, 1 for C
ZIP=37216
active_hysteresis = 0.0
inactive_hysteresis = 0.5
outdoor_temp_buffer = 5

#GPIO pins
HEATER_PIN = 23
AC_PIN = 24
FAN_PIN = 25

#NOTE: thermometer will be read via Dallas 1-wire, as per
#http://learn.adafruit.com/adafruits-raspberry-pi-lesson-11-ds18b20-temperature-sensing/software

#Setting up GPIO
GPIO.setmode(GPIO.BCM)
GPIO.setup(HEATER_PIN, GPIO.OUT)
GPIO.setup(AC_PIN, GPIO.OUT)
GPIO.setup(FAN_PIN, GPIO.OUT)

###Begin helper functions###

def getHVACState():
    heatStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(HEATER_PIN) + "/value").read().strip())
    coolStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(AC_PIN) + "/value").read().strip())
    fanStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(FAN_PIN) + "/value").read().strip())
    if heatStatus == 1 and fanStatus == 1:
        #heating
        return 1
    elif coolStatus == 1 and fanStatus == 1:
        #cooling
        return -1
    elif heatStatus == 0 and coolStatus == 0 and fanStatus == 0:
        #idle
        return 0
    else:
        #broken
        return 2

def cool():
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, True)
    GPIO.output(FAN_PIN, True)
    return -1

def heat():
    GPIO.output(HEATER_PIN, True)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, True)
    return 1

def fan_to_idle():
    #to blow the rest of the heated / cooled air out of the system
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, True)

def idle():
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, False)
    return 0

###begin main loop###
#infinite loop is the same as a daemon, right?
while 1 == 1:

    indoor_temp = float(getIndoorTemp())
    hvac_state = int(getHVACState())

    file = open("status", "r")
    set_temp = float(file.readline())
    mode = file.readline()
    file.close()

    # heater mode
    if mode == "heat":
        if hvac_state == 0: #idle
            if indoor_temp < set_temp - inactive_hysteresis:
                hvac_state = heat()
        elif hvac_state == 1: #heating
            if indoor_temp > set_temp + active_hysteresis:
                fan_to_idle()
                time.sleep(30)
                hvac_state = idle()
        elif hvac_state == -1: # it's cold out, why is the AC running?
                hvac_state = idle()
    # ac mode
    elif mode == "cool":
        if hvac_state == 0: #idle
            if indoor_temp > set_temp + inactive_hysteresis:
                hvac_state = cool()
        elif hvac_state == -1: #cooling
            if indoor_temp < set_temp - active_hysteresis:
                fan_to_idle()
                time.sleep(30)
                hvac_state = idle()
        elif hvac_state == 1: # it's hot out, why is the heater on?
                hvac_state = idle()
    else:
        print "It broke."

    #debug stuff
    if DEBUG == 1:
        print "Sleepy time!"
        print "hvac_state = " + str(hvac_state)
        print "indoor_temp = " + str(indoor_temp)
        print "outdoor_temp = " + str(outdoor_temp)
        print "set_temp = " + str(set_temp)

    time.sleep(10)
<commit_msg>remove outdated debug message, set debug on by default<commit_after>#! /usr/bin/python

import sys
import subprocess
import os
import time
from getIndoorTemp import getIndoorTemp
import RPi.GPIO as GPIO

DEBUG = 1

METRIC=0 # 0 for F, 1 for C
ZIP=37216
active_hysteresis = 0.0
inactive_hysteresis = 0.5
outdoor_temp_buffer = 5

#GPIO pins
HEATER_PIN = 23
AC_PIN = 24
FAN_PIN = 25

#NOTE: thermometer will be read via Dallas 1-wire, as per
#http://learn.adafruit.com/adafruits-raspberry-pi-lesson-11-ds18b20-temperature-sensing/software

#Setting up GPIO
GPIO.setmode(GPIO.BCM)
GPIO.setup(HEATER_PIN, GPIO.OUT)
GPIO.setup(AC_PIN, GPIO.OUT)
GPIO.setup(FAN_PIN, GPIO.OUT)

###Begin helper functions###

def getHVACState():
    heatStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(HEATER_PIN) + "/value").read().strip())
    coolStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(AC_PIN) + "/value").read().strip())
    fanStatus = int(os.popen("cat /sys/class/gpio/gpio" + str(FAN_PIN) + "/value").read().strip())
    if heatStatus == 1 and fanStatus == 1:
        #heating
        return 1
    elif coolStatus == 1 and fanStatus == 1:
        #cooling
        return -1
    elif heatStatus == 0 and coolStatus == 0 and fanStatus == 0:
        #idle
        return 0
    else:
        #broken
        return 2

def cool():
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, True)
    GPIO.output(FAN_PIN, True)
    return -1

def heat():
    GPIO.output(HEATER_PIN, True)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, True)
    return 1

def fan_to_idle():
    #to blow the rest of the heated / cooled air out of the system
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, True)

def idle():
    GPIO.output(HEATER_PIN, False)
    GPIO.output(AC_PIN, False)
    GPIO.output(FAN_PIN, False)
    return 0

###begin main loop###
#infinite loop is the same as a daemon, right?
while 1 == 1:

    indoor_temp = float(getIndoorTemp())
    hvac_state = int(getHVACState())

    file = open("status", "r")
    set_temp = float(file.readline())
    mode = file.readline()
    file.close()

    # heater mode
    if mode == "heat":
        if hvac_state == 0: #idle
            if indoor_temp < set_temp - inactive_hysteresis:
                hvac_state = heat()
        elif hvac_state == 1: #heating
            if indoor_temp > set_temp + active_hysteresis:
                fan_to_idle()
                time.sleep(30)
                hvac_state = idle()
        elif hvac_state == -1: # it's cold out, why is the AC running?
                hvac_state = idle()
    # ac mode
    elif mode == "cool":
        if hvac_state == 0: #idle
            if indoor_temp > set_temp + inactive_hysteresis:
                hvac_state = cool()
        elif hvac_state == -1: #cooling
            if indoor_temp < set_temp - active_hysteresis:
                fan_to_idle()
                time.sleep(30)
                hvac_state = idle()
        elif hvac_state == 1: # it's hot out, why is the heater on?
                hvac_state = idle()
    else:
        print "It broke."

    #debug stuff
    if DEBUG == 1:
        print "Sleepy time!"
        print "hvac_state = " + str(hvac_state)
        print "indoor_temp = " + str(indoor_temp)
        print "set_temp = " + str(set_temp)

    time.sleep(10)
<|endoftext|>


<commit_before>"""Compatible Relaxation"""
from __future__ import print_function

__docformat__ = "restructuredtext en"

import numpy as np
import scipy as sp
from scipy.linalg import norm
from scipy.sparse import isspmatrix, spdiags, isspmatrix_csr

from ..relaxation.relaxation import gauss_seidel, gauss_seidel_indexed

__all__ = ['CR', 'binormalize']

def CR(S, method='habituated', maxiter=20):
    """Use Compatible Relaxation to compute a C/F splitting

    Parameters
    ----------
    S : csr_matrix
        sparse matrix (n x n) usually matrix A of Ax=b
    method : {'habituated','concurrent'}
        Method used during relaxation:
            - concurrent: GS relaxation on F-points, leaving e_c = 0
            - habituated: full relaxation, setting e_c = 0
    maxiter : int
        maximum number of outer iterations (lambda)

    Returns
    -------
    splitting : array
        C/F list of 1's (coarse pt) and 0's (fine pt) (n x 1)

    References
    ----------
    .. [1] Livne, O.E., "Coarsening by compatible relaxation."
       Numer. Linear Algebra Appl. 11, No. 2-3, 205-227 (2004).

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.classical.cr import CR
    >>> A = poisson((20,20),format='csr')
    >>> splitting = CR(A)

    """
    # parameters (paper notation)
    ntests = 3      # (nu) number of random tests to do per iteration
    nrelax = 4      # (eta) number of relaxation sweeps per test

    smagic = 1.0    # (s) parameter in [1,5] to account for fill-in
    gamma = 1.5     # (gamma) cycle index.  use 1.5 for 2d
    G = 30          # (G) number of equivalence classes (# of bins)
    tdepth = 1      # (t) drop depth on parse of L bins
    delta = 0       # (delta) drop threshold on parse of L bins
    alphai = 0.25   # (alpha_inc) quota increase

    # initializations
    alpha = 0.0     # coarsening ratio, quota
    beta = np.inf      # quality criterion
    beta1 = np.inf     # quality criterion, older
    beta2 = np.inf     # quality criterion, oldest
    n = S.shape[0]    # problem size
    nC = 0          # number of current Coarse points
    # rhs = np.zeros((n, 1))  # rhs for Ae=0

    if not isspmatrix(S):
        raise TypeError('expecting sparse matrix')

    S = binormalize(S)

    splitting = np.zeros((S.shape[0], 1), dtype='intc')

    # out iterations ---------------
    for m in range(0, maxiter):

        mu = 0.0  # convergence rate
        E = np.zeros((n, 1))  # slowness measure

        # random iterations ---------------
        for k in range(0, ntests):

            e = 0.5*(1 + sp.rand(n, 1))
            e[splitting > 0] = 0

            enorm = norm(e)

            # relaxation iterations ---------------
            for l in range(0, nrelax):

                if method == 'habituated':
                    gauss_seidel(S, e, np.zeros((n, 1)), iterations=1)
                    e[splitting > 0] = 0
                elif method == 'concurrent':
                    raise NotImplementedError('not implemented: need an \
                                                F-smoother')
                else:
                    raise NotImplementedError('method not recognized: need \
                                               habituated or concurrent')

                enorm_old = enorm
                enorm = norm(e)

                if enorm <= 1e-14:
                    # break out of loops
                    ntests = k
                    nrelax = l
                    maxiter = m
            # end relax

            # check slowness
            E = np.where(np.abs(e) > E, np.abs(e), E)

            # update convergence rate
            mu = mu + enorm/enorm_old
        # end random tests
        mu = mu/ntests

        # work
        alpha = float(nC)/n

        W = (1 + (smagic-1)*gamma*alpha)/(1-gamma*alpha)

        # quality criterion
        beta2 = beta1
        beta1 = beta
        beta = np.power(max([mu, 0.1]), 1.0 / W)

        # check if we're doing well
        if (beta > beta1 and beta1 > beta2) or \
           m == (maxiter-1) or max(E) < 1e-13:
            return splitting.ravel()

        # now add points
        #
        # update limit on additions to splitting (C)
        if alpha < 1e-13:
            alpha = 0.25
        else:
            alpha = (1-alphai) * alpha + alphai * (1/gamma)

        nCmax = np.ceil(alpha * n)

        L = np.ceil(G * E / E.max()).ravel()

        binid = G

        # add whole bins (and t-depth nodes) at a time
        # u = np.zeros((n, 1))
        # TODO This loop may never halt...
        #      Perhaps loop over nC < nCmax and binid > 0 ?
        while nC < nCmax:
            if delta > 0:
                raise NotImplementedError
            if tdepth != 1:
                raise NotImplementedError

            (roots,) = np.where(L == binid)

            for root in roots:
                if L[root] >= 0:
                    cols = S[root, :].indices
                    splitting[root] = 1    # add roots
                    nC += 1
                    L[cols] = -1
            binid -= 1

            # L[troots] = -1          # mark t-rings visited
            # u[:]=0.0
            # u[roots] = 1.0
            # for depth in range(0,tdepth):
            #     u = np.abs(S) * u
            # (troots,tmp) = np.where(u>0)

    return splitting.ravel()

def _CRsweep(A, Findex, Cindex, nu, thetacr, method):

    n = A.shape[0]    # problem size
    numax = nu
    z = np.zeros((n,))
    e = np.ones((n,))
    e[Cindex] = 0.0
    enorm = norm(e)
    rhok = 1

    for it in range(1, numax+1):

        if method == 'habituated':
            gauss_seidel(A, e, z, iterations=1)
            e[Cindex] = 0.0
        elif method == 'concurrent':
            gauss_seidel_indexed(A, e, z, indices=Findex, iterations=1)
        else:
            raise NotImplementedError('method not recognized: need \
                                       habituated or concurrent')

        enorm_old = enorm
        enorm = norm(e)
        rhok_old = rhok
        rhok = enorm / enorm_old

        # criteria 1
        if (abs(rhok - rhok_old) / rhok < 0.1) and (it >= nu):
            return rhok, e

        # criteria 2
        if rhok < 0.1 * thetacr:
            return rhok, e

def CRalpha(A, method='habituated', nu=3, thetacr=0.7, thetacs=[0.3, 0.5],
            maxiter=20):
    """
    >>> from pyamg.gallery import poisson
    >>> from cr import CRalpha
    >>> A = poisson((20,20),format='csr')
    >>> splitting = CRalpha(A)
    """
    n = A.shape[0]    # problem size

    thetacs = list(thetacs)
    thetacs.reverse()

    if not isspmatrix_csr(A):
        raise TypeError('expecting csr sparse matrix A')

    if A.dtype == complex:
        raise NotImplementedError('complex A not implemented')

    # 3.1a
    splitting = np.zeros((n,), dtype='intc')
    gamma = np.zeros((n,))

    # 3.1b
    Cindex = np.where(splitting == 1)[0]
    Findex = np.where(splitting == 0)[0]
    rho, e = _CRsweep(A, Findex, Cindex, nu, thetacr, method=method)

    # 3.1c
    for it in range(0, maxiter):

        print(it)
        # 3.1d (assuming constant initial e in _CRsweep)
        # should already be zero at C pts (Cindex)
        gamma[Findex] = np.abs(e[Findex]) / np.abs(e[Findex]).max()

        # 3.1e
        Uindex = np.where(gamma > thetacs[0])[0]
        if len(thetacs) > 1:
            thetacs.pop()

        # 3.1f
        # first find the weights: omega_i = |N_i\C| + gamma_i
        omega = -np.inf * np.ones((n,))
        for i in Uindex:
            J = A.indices[np.arange(A.indptr[i], A.indptr[i+1])]
            J = np.where(splitting[J] == 0)[0]
            omega[i] = len(J) + gamma[i]

        # independent set
        Usize = len(Uindex)
        while Usize > 0:
            # step 1
            i = omega.argmax()
            splitting[i] = 1
            gamma[i] = 0.0
            # step 2
            J = A.indices[np.arange(A.indptr[i], A.indptr[i+1])]
            J = np.intersect1d(J, Uindex, assume_unique=True)
            omega[i] = -np.inf
            omega[J] = -np.inf

            # step 3
            for j in J:
                K = A.indices[np.arange(A.indptr[j], A.indptr[j+1])]
                K = np.intersect1d(K, Uindex, assume_unique=True)
                omega[K] = omega[K] + 1.0

            Usize -= 1

        Cindex = np.where(splitting == 1)[0]
        Findex = np.where(splitting == 0)[0]
        rho, e = _CRsweep(A, Findex, Cindex, nu, thetacr, method=method)

        print(rho)
        if rho < thetacr:
            break

    return splitting

def binormalize(A, tol=1e-5, maxiter=10):
    """Binormalize matrix A.  Attempt to create unit l_1 norm rows.

    Parameters
    ----------
    A : csr_matrix
        sparse matrix (n x n)
    tol : float
        tolerance
    x : array
        guess at the diagonal
    maxiter : int
        maximum number of iterations to try

    Returns
    -------
    C : csr_matrix
        diagonally scaled A, C=DAD

    Notes
    -----
        - Goal: Scale A so that l_1 norm of the rows are equal to 1:
        - B = DAD
        - want row sum of B = 1
        - easily done with tol=0 if B=DA, but this is not symmetric
        - algorithm is O(N log (1.0/tol))

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.classical import binormalize
    >>> A = poisson((10,),format='csr')
    >>> C = binormalize(A)

    References
    ----------
    .. [1] Livne, Golub, "Scaling by Binormalization"
       Tech Report SCCM-03-12, SCCM, Stanford, 2003
       http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.3.1679

    """
    if not isspmatrix(A):
        raise TypeError('expecting sparse matrix A')

    if A.dtype == complex:
        raise NotImplementedError('complex A not implemented')

    n = A.shape[0]
    it = 0
    x = np.ones((n, 1)).ravel()

    # 1.
    B = A.multiply(A).tocsc()  # power(A,2) inconsistent in numpy, scipy.sparse
    d = B.diagonal().ravel()

    # 2.
    beta = B * x
    betabar = (1.0/n) * np.dot(x, beta)
    stdev = rowsum_stdev(x, beta)

    # 3
    while stdev > tol and it < maxiter:
        for i in range(0, n):
            # solve equation x_i, keeping x_j's fixed
            # see equation (12)
            c2 = (n-1)*d[i]
            c1 = (n-2)*(beta[i] - d[i]*x[i])
            c0 = -d[i]*x[i]*x[i] + 2*beta[i]*x[i] - n*betabar
            if (-c0 < 1e-14):
                print('warning: A nearly un-binormalizable...')
                return A
            else:
                # see equation (12)
                xnew = (2*c0)/(-c1 - np.sqrt(c1*c1 - 4*c0*c2))
            dx = xnew - x[i]

            # here we assume input matrix is symmetric since we grab a row of B
            # instead of a column
            ii = B.indptr[i]
            iii = B.indptr[i+1]
            dot_Bcol = np.dot(x[B.indices[ii:iii]], B.data[ii:iii])

            betabar = betabar + (1.0/n)*dx*(dot_Bcol + beta[i] + d[i]*dx)
            beta[B.indices[ii:iii]] += dx*B.data[ii:iii]

            x[i] = xnew

        stdev = rowsum_stdev(x, beta)
        it += 1

    # rescale for unit 2-norm
    d = np.sqrt(x)
    D = spdiags(d.ravel(), [0], n, n)
    C = D * A * D
    C = C.tocsr()
    beta = C.multiply(C).sum(axis=1)
    scale = np.sqrt((1.0/n) * np.sum(beta))
    return (1/scale)*C

def rowsum_stdev(x, beta):
    """Compute row sum standard deviation

    Compute for approximation x, the std dev of the row sums
    s(x) = ( 1/n \sum_k  (x_k beta_k - betabar)^2 )^(1/2)
    with betabar = 1/n dot(beta,x)

    Parameters
    ----------
    x : array
    beta : array

    Returns
    -------
    s(x)/betabar : float

    Notes
    -----
    equation (7) in Livne/Golub

    """
    n = x.size
    betabar = (1.0/n) * np.dot(x, beta)
    stdev = np.sqrt((1.0/n) *
                    np.sum(np.power(np.multiply(x, beta) - betabar, 2)))
    return stdev/betabar
<commit_msg>added concurrent cr<commit_after>"""Compatible Relaxation"""
from __future__ import print_function

__docformat__ = "restructuredtext en"

import numpy as np
import scipy as sp
from scipy.linalg import norm
from scipy.sparse import isspmatrix, spdiags, isspmatrix_csr

from ..relaxation.relaxation import gauss_seidel, gauss_seidel_indexed

__all__ = ['CR', 'binormalize']

def CR(S, method='habituated', maxiter=20):
    """Use Compatible Relaxation to compute a C/F splitting

    Parameters
    ----------
    S : csr_matrix
        sparse matrix (n x n) usually matrix A of Ax=b
    method : {'habituated','concurrent'}
        Method used during relaxation:
            - concurrent: GS relaxation on F-points, leaving e_c = 0
            - habituated: full relaxation, setting e_c = 0
    maxiter : int
        maximum number of outer iterations (lambda)

    Returns
    -------
    splitting : array
        C/F list of 1's (coarse pt) and 0's (fine pt) (n x 1)

    References
    ----------
    .. [1] Livne, O.E., "Coarsening by compatible relaxation."
       Numer. Linear Algebra Appl. 11, No. 2-3, 205-227 (2004).

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.classical.cr import CR
    >>> A = poisson((20,20),format='csr')
    >>> splitting = CR(A)

    """
    # parameters (paper notation)
    ntests = 3      # (nu) number of random tests to do per iteration
    nrelax = 4      # (eta) number of relaxation sweeps per test

    smagic = 1.0    # (s) parameter in [1,5] to account for fill-in
    gamma = 1.5     # (gamma) cycle index.  use 1.5 for 2d
    G = 30          # (G) number of equivalence classes (# of bins)
    tdepth = 1      # (t) drop depth on parse of L bins
    delta = 0       # (delta) drop threshold on parse of L bins
    alphai = 0.25   # (alpha_inc) quota increase

    # initializations
    alpha = 0.0     # coarsening ratio, quota
    beta = np.inf      # quality criterion
    beta1 = np.inf     # quality criterion, older
    beta2 = np.inf     # quality criterion, oldest
    n = S.shape[0]    # problem size
    nC = 0          # number of current Coarse points
    rhs = np.zeros((n, 1))  # rhs for Ae=0

    if not isspmatrix(S):
        raise TypeError('expecting sparse matrix')

    S = binormalize(S)

    splitting = np.zeros((S.shape[0], 1), dtype='intc')

    # out iterations ---------------
    for m in range(0, maxiter):

        Cpts = np.where(splitting == 1)[0]
        Fpts = np.where(splitting == 0)[0]
        mu = 0.0  # convergence rate
        E = np.zeros((n, 1))  # slowness measure

        # random iterations ---------------
        for k in range(0, ntests):

            e = 0.5*(1 + sp.rand(n, 1))
            e[Cpts] = 0

            enorm = norm(e)

            # relaxation iterations ---------------
            for l in range(0, nrelax):

                if method == 'habituated':
                    gauss_seidel(S, e, rhs, iterations=1)
                    e[Cpts] = 0
                elif method == 'concurrent':
                    gauss_seidel_indexed(S, e, rhs, indices=Fpts, iterations=1)
                else:
                    raise NotImplementedError('method not recognized: need \
                                               habituated or concurrent')

                enorm_old = enorm
                enorm = norm(e)

                if enorm <= 1e-14:
                    # break out of loops
                    ntests = k
                    nrelax = l
                    maxiter = m
            # end relax

            # check slowness
            E = np.where(np.abs(e) > E, np.abs(e), E)

            # update convergence rate
            mu = mu + enorm/enorm_old
        # end random tests
        mu = mu/ntests

        # work
        alpha = float(nC)/n

        W = (1 + (smagic-1)*gamma*alpha)/(1-gamma*alpha)

        # quality criterion
        beta2 = beta1
        beta1 = beta
        beta = np.power(max([mu, 0.1]), 1.0 / W)

        # check if we're doing well
        if (beta > beta1 and beta1 > beta2) or \
           m == (maxiter-1) or max(E) < 1e-13:
            return splitting.ravel()

        # now add points
        #
        # update limit on additions to splitting (C)
        if alpha < 1e-13:
            alpha = 0.25
        else:
            alpha = (1-alphai) * alpha + alphai * (1/gamma)

        nCmax = np.ceil(alpha * n)

        L = np.ceil(G * E / E.max()).ravel()

        binid = G

        # add whole bins (and t-depth nodes) at a time
        # u = np.zeros((n, 1))
        # TODO This loop may never halt...
        #      Perhaps loop over nC < nCmax and binid > 0 ?
        while nC < nCmax:
            if delta > 0:
                raise NotImplementedError
            if tdepth != 1:
                raise NotImplementedError

            (roots,) = np.where(L == binid)

            for root in roots:
                if L[root] >= 0:
                    cols = S[root, :].indices
                    splitting[root] = 1    # add roots
                    nC += 1
                    L[cols] = -1
            binid -= 1

            # L[troots] = -1          # mark t-rings visited
            # u[:]=0.0
            # u[roots] = 1.0
            # for depth in range(0,tdepth):
            #     u = np.abs(S) * u
            # (troots,tmp) = np.where(u>0)

    return splitting.ravel()

def _CRsweep(A, Findex, Cindex, nu, thetacr, method):

    n = A.shape[0]    # problem size
    numax = nu
    z = np.zeros((n,))
    e = np.ones((n,))
    e[Cindex] = 0.0
    enorm = norm(e)
    rhok = 1

    for it in range(1, numax+1):

        if method == 'habituated':
            gauss_seidel(A, e, z, iterations=1)
            e[Cindex] = 0.0
        elif method == 'concurrent':
            gauss_seidel_indexed(A, e, z, indices=Findex, iterations=1)
        else:
            raise NotImplementedError('method not recognized: need \
                                       habituated or concurrent')

        enorm_old = enorm
        enorm = norm(e)
        rhok_old = rhok
        rhok = enorm / enorm_old

        # criteria 1
        if (abs(rhok - rhok_old) / rhok < 0.1) and (it >= nu):
            return rhok, e

        # criteria 2
        if rhok < 0.1 * thetacr:
            return rhok, e

def CRalpha(A, method='habituated', nu=3, thetacr=0.7, thetacs=[0.3, 0.5],
            maxiter=20):
    """
    >>> from pyamg.gallery import poisson
    >>> from cr import CRalpha
    >>> A = poisson((20,20),format='csr')
    >>> splitting = CRalpha(A)
    """
    n = A.shape[0]    # problem size

    thetacs = list(thetacs)
    thetacs.reverse()

    if not isspmatrix_csr(A):
        raise TypeError('expecting csr sparse matrix A')

    if A.dtype == complex:
        raise NotImplementedError('complex A not implemented')

    # 3.1a
    splitting = np.zeros((n,), dtype='intc')
    gamma = np.zeros((n,))

    # 3.1b
    Cindex = np.where(splitting == 1)[0]
    Findex = np.where(splitting == 0)[0]
    rho, e = _CRsweep(A, Findex, Cindex, nu, thetacr, method=method)

    # 3.1c
    for it in range(0, maxiter):

        print(it)
        # 3.1d (assuming constant initial e in _CRsweep)
        # should already be zero at C pts (Cindex)
        gamma[Findex] = np.abs(e[Findex]) / np.abs(e[Findex]).max()

        # 3.1e
        Uindex = np.where(gamma > thetacs[0])[0]
        if len(thetacs) > 1:
            thetacs.pop()

        # 3.1f
        # first find the weights: omega_i = |N_i\C| + gamma_i
        omega = -np.inf * np.ones((n,))
        for i in Uindex:
            J = A.indices[np.arange(A.indptr[i], A.indptr[i+1])]
            J = np.where(splitting[J] == 0)[0]
            omega[i] = len(J) + gamma[i]

        # independent set
        Usize = len(Uindex)
        while Usize > 0:
            # step 1
            i = omega.argmax()
            splitting[i] = 1
            gamma[i] = 0.0
            # step 2
            J = A.indices[np.arange(A.indptr[i], A.indptr[i+1])]
            J = np.intersect1d(J, Uindex, assume_unique=True)
            omega[i] = -np.inf
            omega[J] = -np.inf

            # step 3
            for j in J:
                K = A.indices[np.arange(A.indptr[j], A.indptr[j+1])]
                K = np.intersect1d(K, Uindex, assume_unique=True)
                omega[K] = omega[K] + 1.0

            Usize -= 1

        Cindex = np.where(splitting == 1)[0]
        Findex = np.where(splitting == 0)[0]
        rho, e = _CRsweep(A, Findex, Cindex, nu, thetacr, method=method)

        print(rho)
        if rho < thetacr:
            break

    return splitting

def binormalize(A, tol=1e-5, maxiter=10):
    """Binormalize matrix A.  Attempt to create unit l_1 norm rows.

    Parameters
    ----------
    A : csr_matrix
        sparse matrix (n x n)
    tol : float
        tolerance
    x : array
        guess at the diagonal
    maxiter : int
        maximum number of iterations to try

    Returns
    -------
    C : csr_matrix
        diagonally scaled A, C=DAD

    Notes
    -----
        - Goal: Scale A so that l_1 norm of the rows are equal to 1:
        - B = DAD
        - want row sum of B = 1
        - easily done with tol=0 if B=DA, but this is not symmetric
        - algorithm is O(N log (1.0/tol))

    Examples
    --------
    >>> from pyamg.gallery import poisson
    >>> from pyamg.classical import binormalize
    >>> A = poisson((10,),format='csr')
    >>> C = binormalize(A)

    References
    ----------
    .. [1] Livne, Golub, "Scaling by Binormalization"
       Tech Report SCCM-03-12, SCCM, Stanford, 2003
       http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.3.1679

    """
    if not isspmatrix(A):
        raise TypeError('expecting sparse matrix A')

    if A.dtype == complex:
        raise NotImplementedError('complex A not implemented')

    n = A.shape[0]
    it = 0
    x = np.ones((n, 1)).ravel()

    # 1.
    B = A.multiply(A).tocsc()  # power(A,2) inconsistent in numpy, scipy.sparse
    d = B.diagonal().ravel()

    # 2.
    beta = B * x
    betabar = (1.0/n) * np.dot(x, beta)
    stdev = rowsum_stdev(x, beta)

    # 3
    while stdev > tol and it < maxiter:
        for i in range(0, n):
            # solve equation x_i, keeping x_j's fixed
            # see equation (12)
            c2 = (n-1)*d[i]
            c1 = (n-2)*(beta[i] - d[i]*x[i])
            c0 = -d[i]*x[i]*x[i] + 2*beta[i]*x[i] - n*betabar
            if (-c0 < 1e-14):
                print('warning: A nearly un-binormalizable...')
                return A
            else:
                # see equation (12)
                xnew = (2*c0)/(-c1 - np.sqrt(c1*c1 - 4*c0*c2))
            dx = xnew - x[i]

            # here we assume input matrix is symmetric since we grab a row of B
            # instead of a column
            ii = B.indptr[i]
            iii = B.indptr[i+1]
            dot_Bcol = np.dot(x[B.indices[ii:iii]], B.data[ii:iii])

            betabar = betabar + (1.0/n)*dx*(dot_Bcol + beta[i] + d[i]*dx)
            beta[B.indices[ii:iii]] += dx*B.data[ii:iii]

            x[i] = xnew

        stdev = rowsum_stdev(x, beta)
        it += 1

    # rescale for unit 2-norm
    d = np.sqrt(x)
    D = spdiags(d.ravel(), [0], n, n)
    C = D * A * D
    C = C.tocsr()
    beta = C.multiply(C).sum(axis=1)
    scale = np.sqrt((1.0/n) * np.sum(beta))
    return (1/scale)*C

def rowsum_stdev(x, beta):
    """Compute row sum standard deviation

    Compute for approximation x, the std dev of the row sums
    s(x) = ( 1/n \sum_k  (x_k beta_k - betabar)^2 )^(1/2)
    with betabar = 1/n dot(beta,x)

    Parameters
    ----------
    x : array
    beta : array

    Returns
    -------
    s(x)/betabar : float

    Notes
    -----
    equation (7) in Livne/Golub

    """
    n = x.size
    betabar = (1.0/n) * np.dot(x, beta)
    stdev = np.sqrt((1.0/n) *
                    np.sum(np.power(np.multiply(x, beta) - betabar, 2)))
    return stdev/betabar
<|endoftext|>


<commit_before>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re
import string

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            ]
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = string.split(ctx.opts.tree_path, '/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(modules, fd, ctx.opts.tree_depth, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts>   <type>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf or presence container
    *  for a leaf-list
    [<keys>] for a list's keys

  <type> is the name of the type for leafs and leaf-lists
""")

def emit_tree(modules, fd, depth, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                print_header = True
            print_children(chs, module, fd, ' ', path, depth)

        rpcs = module.search('rpc')
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                print_header = True
            fd.write("rpcs:\n")
            print_children(rpcs, module, fd, ' ', path, depth)

        notifs = module.search('notification')
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                print_header = True
            fd.write("notifications:\n")
            print_children(notifs, module, fd, ' ', path, depth)

def print_children(i_children, module, fd, prefix, path, depth, width=0):
    if depth == 0:
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ch == i_children[-1]:
            newprefix = prefix + '   '
        else:
            newprefix = prefix + '  |'
        if ((ch.arg == 'input' or ch.arg == 'output') and
            ch.parent.keyword == 'rpc' and
            len(ch.i_children) == 0 and
            ch.parent.search_one(ch.arg) is None):
            pass
        else:
            print_node(ch, module, fd, newprefix, path, depth, width)

def print_node(s, module, fd, prefix, path, depth, width):
    fd.write("%s%s--" % (prefix[0:-1], get_status_str(s)))

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s)
    if s.keyword == 'list':
        fd.write(flags + " " + name)
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '?'
        fd.write(flags + " " + name)
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            fd.write(flags + ' (' + s.arg + ')?')
        else:
            fd.write(flags + ' (' + s.arg + ')')
    elif s.keyword == 'case':
        fd.write(':(' + s.arg + ')')
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif s.keyword == 'leaf' and not hasattr(s, 'i_is_key'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        fd.write("%s %-*s   %s" % (flags, width+1, name, get_typename(s)))

    if s.keyword == 'list' and s.search_one('key') is not None:
        fd.write(" [%s]" % re.sub('\s+', ' ', s.search_one('key').arg))
    fd.write('\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, depth, width)
        else:
            print_children(chs, module, fd, prefix, path, depth)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s):
    if s.keyword == 'rpc':
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    else:
        return 'ro'

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        return t.arg
    else:
        return ''
<commit_msg>Fix typo.<commit_after>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re
import string

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            ]
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = string.split(ctx.opts.tree_path, '/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(modules, fd, ctx.opts.tree_depth, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts>   <type>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf or presence container
    *  for a leaf-list
    [<keys>] for a list's keys

  <type> is the name of the type for leafs and leaf-lists
""")

def emit_tree(modules, fd, depth, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            print_children(chs, module, fd, ' ', path, depth)

        rpcs = module.search('rpc')
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("rpcs:\n")
            print_children(rpcs, module, fd, ' ', path, depth)

        notifs = module.search('notification')
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("notifications:\n")
            print_children(notifs, module, fd, ' ', path, depth)

def print_children(i_children, module, fd, prefix, path, depth, width=0):
    if depth == 0:
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ch == i_children[-1]:
            newprefix = prefix + '   '
        else:
            newprefix = prefix + '  |'
        if ((ch.arg == 'input' or ch.arg == 'output') and
            ch.parent.keyword == 'rpc' and
            len(ch.i_children) == 0 and
            ch.parent.search_one(ch.arg) is None):
            pass
        else:
            print_node(ch, module, fd, newprefix, path, depth, width)

def print_node(s, module, fd, prefix, path, depth, width):
    fd.write("%s%s--" % (prefix[0:-1], get_status_str(s)))

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s)
    if s.keyword == 'list':
        fd.write(flags + " " + name)
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '?'
        fd.write(flags + " " + name)
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            fd.write(flags + ' (' + s.arg + ')?')
        else:
            fd.write(flags + ' (' + s.arg + ')')
    elif s.keyword == 'case':
        fd.write(':(' + s.arg + ')')
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif s.keyword == 'leaf' and not hasattr(s, 'i_is_key'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        fd.write("%s %-*s   %s" % (flags, width+1, name, get_typename(s)))

    if s.keyword == 'list' and s.search_one('key') is not None:
        fd.write(" [%s]" % re.sub('\s+', ' ', s.search_one('key').arg))
    fd.write('\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, depth, width)
        else:
            print_children(chs, module, fd, prefix, path, depth)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s):
    if s.keyword == 'rpc':
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    else:
        return 'ro'

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        return t.arg
    else:
        return ''
<|endoftext|>


<commit_before>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re
import string

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            ]
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = string.split(ctx.opts.tree_path, '/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(ctx, modules, fd, ctx.opts.tree_depth, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts> <type> <if-features>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf or choice
    !  for a presence container
    *  for a leaf-list or list
    [<keys>] for a list's keys

  <type> is the name of the type for leafs and leaf-lists

    If the type is a leafref, the type is printed as "-> TARGET", where
    TARGET is either the leafref path, with prefixed removed if possible.

  <if-features> is the list of features this node depends on, printed
    within curly brackets and a question mark "{...}?"
""")

def emit_tree(ctx, modules, fd, depth, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))
            printed_header = True

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            print_children(chs, module, fd, ' ', path, 'data', depth)

        mods = [module]
        for i in module.search('include'):
            subm = ctx.get_module(i.arg)
            if subm is not None:
                mods.append(subm)
        for m in mods:
            for augment in m.search('augment'):
                if (hasattr(augment.i_target_node, 'i_module') and
                    augment.i_target_node.i_module not in modules + mods):
                    # this augment has not been printed; print it
                    if not printed_header:
                        print_header()
                        printed_header = True
                    fd.write("augment %s:\n" % augment.arg)
                    print_children(augment.i_children, m, fd,
                                   ' ', path, 'augment', depth)

        rpcs = [ch for ch in module.i_children
                if ch.keyword == 'rpc']
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("rpcs:\n")
            print_children(rpcs, module, fd, ' ', path, 'rpc', depth)

        notifs = [ch for ch in module.i_children
                  if ch.keyword == 'notification']
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("notifications:\n")
            print_children(notifs, module, fd, ' ', path, 'notification', depth)

def print_children(i_children, module, fd, prefix, path, mode, depth, width=0):
    if depth == 0:
        if i_children: fd.write(prefix + '     ...\n')
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ((ch.keyword == 'input' or ch.keyword == 'output') and
            len(ch.i_children) == 0):
            pass
        else:
            if (ch == i_children[-1] or
                (i_children[-1].keyword == 'output' and
                 len(i_children[-1].i_children) == 0)):
                # the last test is to detect if we print input, and the
                # next node is an empty output node; then don't add the |
                newprefix = prefix + '   '
            else:
                newprefix = prefix + '  |'
            if ch.keyword == 'input':
                mode = 'input'
            elif ch.keyword == 'output':
                mode = 'output'
            print_node(ch, module, fd, newprefix, path, mode, depth, width)

def print_node(s, module, fd, prefix, path, mode, depth, width):
    fd.write("%s%s--" % (prefix[0:-1], get_status_str(s)))

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s, mode)
    if s.keyword == 'list':
        name += '*'
        fd.write(flags + " " + name)
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '!'
        fd.write(flags + " " + name)
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            fd.write(flags + ' (' + s.arg + ')?')
        else:
            fd.write(flags + ' (' + s.arg + ')')
    elif s.keyword == 'case':
        fd.write(':(' + s.arg + ')')
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif s.keyword == 'leaf' and not hasattr(s, 'i_is_key'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        t = get_typename(s)
        if t == '':
            fd.write("%s %s" % (flags, name))
        else:
            fd.write("%s %-*s   %s" % (flags, width+1, name, t))

    if s.keyword == 'list' and s.search_one('key') is not None:
        fd.write(" [%s]" % re.sub('\s+', ' ', s.search_one('key').arg))

    features = s.search('if-feature')
    if len(features) > 0:
        fd.write(" {%s}?" % ",".join([f.arg for f in features]))

    fd.write('\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, mode, depth, width)
        else:
            print_children(chs, module, fd, prefix, path, mode, depth)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s, mode):
    if mode == 'input':
        return "-w"
    elif (s.keyword == 'rpc' or s.keyword == ('tailf-common', 'action')):
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    elif s.i_config == False or mode == 'output' or mode == 'notification':
        return 'ro'
    else:
        return '--'

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            p = t.search_one('path')
            if p is not None:
                # Try to make the path as compact as possible.
                # Remove local prefixes, and only use prefix when
                # there is a module change in the path.
                target = []
                curprefix = s.i_module.i_prefix
                for name in p.arg.split('/'):
                    if name.find(":") == -1:
                        prefix = curprefix
                    else:
                        [prefix, name] = name.split(':', 1)
                    if prefix == curprefix:
                        target.append(name)
                    else:
                        target.append(prefix + ':' + name)
                        curprefix = prefix
                return "-> %s" % "/".join(target)
            else:
                return t.arg
        else:
            return t.arg
    else:
        return ''
<commit_msg>print augmented choices/cases with module prefix in tree output<commit_after>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re
import string

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            ]
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = string.split(ctx.opts.tree_path, '/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(ctx, modules, fd, ctx.opts.tree_depth, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts> <type> <if-features>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf or choice
    !  for a presence container
    *  for a leaf-list or list
    [<keys>] for a list's keys

  <type> is the name of the type for leafs and leaf-lists

    If the type is a leafref, the type is printed as "-> TARGET", where
    TARGET is either the leafref path, with prefixed removed if possible.

  <if-features> is the list of features this node depends on, printed
    within curly brackets and a question mark "{...}?"
""")

def emit_tree(ctx, modules, fd, depth, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))
            printed_header = True

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            print_children(chs, module, fd, ' ', path, 'data', depth)

        mods = [module]
        for i in module.search('include'):
            subm = ctx.get_module(i.arg)
            if subm is not None:
                mods.append(subm)
        for m in mods:
            for augment in m.search('augment'):
                if (hasattr(augment.i_target_node, 'i_module') and
                    augment.i_target_node.i_module not in modules + mods):
                    # this augment has not been printed; print it
                    if not printed_header:
                        print_header()
                        printed_header = True
                    fd.write("augment %s:\n" % augment.arg)
                    print_children(augment.i_children, m, fd,
                                   ' ', path, 'augment', depth)

        rpcs = [ch for ch in module.i_children
                if ch.keyword == 'rpc']
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("rpcs:\n")
            print_children(rpcs, module, fd, ' ', path, 'rpc', depth)

        notifs = [ch for ch in module.i_children
                  if ch.keyword == 'notification']
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("notifications:\n")
            print_children(notifs, module, fd, ' ', path, 'notification', depth)

def print_children(i_children, module, fd, prefix, path, mode, depth, width=0):
    if depth == 0:
        if i_children: fd.write(prefix + '     ...\n')
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ((ch.keyword == 'input' or ch.keyword == 'output') and
            len(ch.i_children) == 0):
            pass
        else:
            if (ch == i_children[-1] or
                (i_children[-1].keyword == 'output' and
                 len(i_children[-1].i_children) == 0)):
                # the last test is to detect if we print input, and the
                # next node is an empty output node; then don't add the |
                newprefix = prefix + '   '
            else:
                newprefix = prefix + '  |'
            if ch.keyword == 'input':
                mode = 'input'
            elif ch.keyword == 'output':
                mode = 'output'
            print_node(ch, module, fd, newprefix, path, mode, depth, width)

def print_node(s, module, fd, prefix, path, mode, depth, width):
    fd.write("%s%s--" % (prefix[0:-1], get_status_str(s)))

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s, mode)
    if s.keyword == 'list':
        name += '*'
        fd.write(flags + " " + name)
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '!'
        fd.write(flags + " " + name)
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            fd.write(flags + ' (' + name + ')?')
        else:
            fd.write(flags + ' (' + name + ')')
    elif s.keyword == 'case':
        fd.write(':(' + name + ')')
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif s.keyword == 'leaf' and not hasattr(s, 'i_is_key'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        t = get_typename(s)
        if t == '':
            fd.write("%s %s" % (flags, name))
        else:
            fd.write("%s %-*s   %s" % (flags, width+1, name, t))

    if s.keyword == 'list' and s.search_one('key') is not None:
        fd.write(" [%s]" % re.sub('\s+', ' ', s.search_one('key').arg))

    features = s.search('if-feature')
    if len(features) > 0:
        fd.write(" {%s}?" % ",".join([f.arg for f in features]))

    fd.write('\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, mode, depth, width)
        else:
            print_children(chs, module, fd, prefix, path, mode, depth)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s, mode):
    if mode == 'input':
        return "-w"
    elif (s.keyword == 'rpc' or s.keyword == ('tailf-common', 'action')):
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    elif s.i_config == False or mode == 'output' or mode == 'notification':
        return 'ro'
    else:
        return '--'

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            p = t.search_one('path')
            if p is not None:
                # Try to make the path as compact as possible.
                # Remove local prefixes, and only use prefix when
                # there is a module change in the path.
                target = []
                curprefix = s.i_module.i_prefix
                for name in p.arg.split('/'):
                    if name.find(":") == -1:
                        prefix = curprefix
                    else:
                        [prefix, name] = name.split(':', 1)
                    if prefix == curprefix:
                        target.append(name)
                    else:
                        target.append(prefix + ':' + name)
                        curprefix = prefix
                return "-> %s" % "/".join(target)
            else:
                return t.arg
        else:
            return t.arg
    else:
        return ''
<|endoftext|>


<commit_before>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def __init__(self):
        plugin.PyangPlugin.__init__(self, 'tree')

    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-line-length",
                                 type="int",
                                 dest="tree_line_length",
                                 help="Maximum line length"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            optparse.make_option("--tree-print-groupings",
                                 dest="tree_print_groupings",
                                 action="store_true",
                                 help="Print groupings"),
            ]
        if plugin.is_plugin_registered('restconf'):
            optlist.append(
                optparse.make_option("--tree-print-yang-data",
                                     dest="tree_print_yang_data",
                                     action="store_true",
                                     help="Print ietf-restconf:yang-data " +
                                     "structures")
            )
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = ctx.opts.tree_path.split('/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(ctx, modules, fd, ctx.opts.tree_depth,
                  ctx.opts.tree_line_length, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts> <type> <if-features>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs and actions
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf, choice, anydata or anyxml
    !  for a presence container
    *  for a leaf-list or list
    [<keys>] for a list's keys

    <type> is the name of the type for leafs and leaf-lists, or
           "<anydata>" or "<anyxml>" for anydata and anyxml, respectively

    If the type is a leafref, the type is printed as "-> TARGET", where
    TARGET is the leafref path, with prefix removed if possible.

  <if-features> is the list of features this node depends on, printed
    within curly brackets and a question mark "{...}?"
""")

def emit_tree(ctx, modules, fd, depth, llen, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))
            printed_header = True

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            print_children(chs, module, fd, '  ', path, 'data', depth, llen)

        mods = [module]
        for i in module.search('include'):
            subm = ctx.get_module(i.arg)
            if subm is not None:
                mods.append(subm)
        for m in mods:
            section_delimiter_printed=False
            for augment in m.search('augment'):
                if (hasattr(augment.i_target_node, 'i_module') and
                    augment.i_target_node.i_module not in modules + mods):
                    if not section_delimiter_printed:
                        fd.write('\n')
                        section_delimiter_printed = True
                    # this augment has not been printed; print it
                    if not printed_header:
                        print_header()
                        printed_header = True
                    fd.write("  augment %s:\n" % augment.arg)
                    print_children(augment.i_children, m, fd,
                                   '  ', path, 'augment', depth, llen)

        rpcs = [ch for ch in module.i_children
                if ch.keyword == 'rpc']
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("\n  rpcs:\n")
            print_children(rpcs, module, fd, '  ', path, 'rpc', depth, llen)

        notifs = [ch for ch in module.i_children
                  if ch.keyword == 'notification']
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("\n  notifications:\n")
            print_children(notifs, module, fd, '  ', path,
                           'notification', depth, llen)

        if ctx.opts.tree_print_groupings and len(module.i_groupings) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            section_delimiter_printed = False
            for gname in module.i_groupings:
                if not section_delimiter_printed:
                    fd.write('\n')
                    section_delimiter_printed = True
                fd.write("  grouping %s\n" % gname)
                g = module.i_groupings[gname]
                print_children(g.i_children, module, fd,
                               '  ', path, 'grouping', depth, llen)

        if ctx.opts.tree_print_yang_data:
            yds = module.search(('ietf-restconf', 'yang-data'))
            if len(yds) > 0:
                if not printed_header:
                    print_header()
                    printed_header = True
                section_delimiter_printed = False
                for yd in yds:
                    if not section_delimiter_printed:
                        fd.write('\n')
                        section_delimiter_printed = True
                    fd.write("  yang-data %s:\n" % yd.arg)
                    print_children(yd.i_children, module, fd, '    ', path,
                                   'yang-data', depth, llen)

def print_children(i_children, module, fd, prefix, path, mode, depth,
                   llen, width=0):
    if depth == 0:
        if i_children: fd.write(prefix + '     ...\n')
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ((ch.keyword == 'input' or ch.keyword == 'output') and
            len(ch.i_children) == 0):
            pass
        else:
            if (ch == i_children[-1] or
                (i_children[-1].keyword == 'output' and
                 len(i_children[-1].i_children) == 0)):
                # the last test is to detect if we print input, and the
                # next node is an empty output node; then don't add the |
                newprefix = prefix + '   '
            else:
                newprefix = prefix + '  |'
            if ch.keyword == 'input':
                mode = 'input'
            elif ch.keyword == 'output':
                mode = 'output'
            print_node(ch, module, fd, newprefix, path, mode, depth, llen,
                       width)

def print_node(s, module, fd, prefix, path, mode, depth, llen, width):
    line = "%s%s--" % (prefix[0:-1], get_status_str(s))

    brcol = len(line) + 4

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s, mode)
    if s.keyword == 'list':
        name += '*'
        line += flags + " " + name
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '!'
        line += flags + " " + name
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            line += flags + ' (' + name + ')?'
        else:
            line += flags + ' (' + name + ')'
    elif s.keyword == 'case':
        line += ':(' + name + ')'
        brcol += 1
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif (s.keyword == 'leaf' and not hasattr(s, 'i_is_key')
              or s.keyword == 'anydata' or s.keyword == 'anyxml'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        t = get_typename(s)
        if t == '':
            line += "%s %s" % (flags, name)
        elif (llen is not None and
              len(line) + len(flags) + width+1 + len(t) + 4 > llen):
            # there's no room for the type name
            if (get_leafref_path(s) is not None and
                len(t) + brcol > llen):
                # it's not even room for the leafref path; skip it
                line += "%s %-*s   leafref" % (flags, width+1, name)
            else:
                line += "%s %s" % (flags, name)
                fd.write(line + '\n')
                line = prefix + ' ' * (brcol - len(prefix)) + ' ' + t
        else:
            line += "%s %-*s   %s" % (flags, width+1, name, t)

    if s.keyword == 'list' and s.search_one('key') is not None:
        keystr = " [%s]" % re.sub('\s+', ' ', s.search_one('key').arg)
        if (llen is not None and
            len(line) + len(keystr) > llen):
            fd.write(line + '\n')
            line = prefix + ' ' * (brcol - len(prefix))
        line += keystr

    features = s.search('if-feature')
    featurenames = [f.arg for f in features]
    if hasattr(s, 'i_augment'):
        afeatures = s.i_augment.search('if-feature')
        featurenames.extend([f.arg for f in afeatures
                             if f.arg not in featurenames])

    if len(featurenames) > 0:
        fstr = " {%s}?" % ",".join(featurenames)
        if (llen is not None and len(line) + len(fstr) > llen):
            fd.write(line + '\n')
            line = prefix + ' ' * (brcol - len(prefix))
        line += fstr

    fd.write(line + '\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, mode, depth,
                           llen, width)
        else:
            print_children(chs, module, fd, prefix, path, mode, depth, llen)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s, mode):
    if mode == 'input':
        return "-w"
    elif s.keyword in ('rpc', 'action', ('tailf-common', 'action')):
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    elif s.i_config == False or mode == 'output' or mode == 'notification':
        return 'ro'
    else:
        return '--'

def get_leafref_path(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            return t.search_one('path')
    else:
        return None

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            p = t.search_one('path')
            if p is not None:
                # Try to make the path as compact as possible.
                # Remove local prefixes, and only use prefix when
                # there is a module change in the path.
                target = []
                curprefix = s.i_module.i_prefix
                for name in p.arg.split('/'):
                    if name.find(":") == -1:
                        prefix = curprefix
                    else:
                        [prefix, name] = name.split(':', 1)
                    if prefix == curprefix:
                        target.append(name)
                    else:
                        target.append(prefix + ':' + name)
                        curprefix = prefix
                return "-> %s" % "/".join(target)
            else:
                return t.arg
        else:
            return t.arg
    elif s.keyword == 'anydata':
        return '<anydata>'
    elif s.keyword == 'anyxml':
        return '<anyxml>'
    else:
        return ''
<commit_msg>typo in help text<commit_after>"""Tree output plugin

Idea copied from libsmi.
"""

import optparse
import sys
import re

from pyang import plugin
from pyang import statements

def pyang_plugin_init():
    plugin.register_plugin(TreePlugin())

class TreePlugin(plugin.PyangPlugin):
    def __init__(self):
        plugin.PyangPlugin.__init__(self, 'tree')

    def add_output_format(self, fmts):
        self.multiple_modules = True
        fmts['tree'] = self

    def add_opts(self, optparser):
        optlist = [
            optparse.make_option("--tree-help",
                                 dest="tree_help",
                                 action="store_true",
                                 help="Print help on tree symbols and exit"),
            optparse.make_option("--tree-depth",
                                 type="int",
                                 dest="tree_depth",
                                 help="Number of levels to print"),
            optparse.make_option("--tree-line-length",
                                 type="int",
                                 dest="tree_line_length",
                                 help="Maximum line length"),
            optparse.make_option("--tree-path",
                                 dest="tree_path",
                                 help="Subtree to print"),
            optparse.make_option("--tree-print-groupings",
                                 dest="tree_print_groupings",
                                 action="store_true",
                                 help="Print groupings"),
            ]
        if plugin.is_plugin_registered('restconf'):
            optlist.append(
                optparse.make_option("--tree-print-yang-data",
                                     dest="tree_print_yang_data",
                                     action="store_true",
                                     help="Print ietf-restconf:yang-data " +
                                     "structures")
            )
        g = optparser.add_option_group("Tree output specific options")
        g.add_options(optlist)

    def setup_ctx(self, ctx):
        if ctx.opts.tree_help:
            print_help()
            sys.exit(0)

    def setup_fmt(self, ctx):
        ctx.implicit_errors = False

    def emit(self, ctx, modules, fd):
        if ctx.opts.tree_path is not None:
            path = ctx.opts.tree_path.split('/')
            if path[0] == '':
                path = path[1:]
        else:
            path = None
        emit_tree(ctx, modules, fd, ctx.opts.tree_depth,
                  ctx.opts.tree_line_length, path)

def print_help():
    print("""
Each node is printed as:

<status> <flags> <name> <opts> <type> <if-features>

  <status> is one of:
    +  for current
    x  for deprecated
    o  for obsolete

  <flags> is one of:
    rw  for configuration data
    ro  for non-configuration data
    -x  for rpcs and actions
    -n  for notifications

  <name> is the name of the node
    (<name>) means that the node is a choice node
   :(<name>) means that the node is a case node

   If the node is augmented into the tree from another module, its
   name is printed as <prefix>:<name>.

  <opts> is one of:
    ?  for an optional leaf, choice, anydata or anyxml
    !  for a presence container
    *  for a leaf-list or list
    [<keys>] for a list's keys

    <type> is the name of the type for leafs and leaf-lists, or
           "<anydata>" or "<anyxml>" for anydata and anyxml, respectively

    If the type is a leafref, the type is printed as "-> TARGET", where
    TARGET is the leafref path, with prefixes removed if possible.

  <if-features> is the list of features this node depends on, printed
    within curly brackets and a question mark "{...}?"
""")

def emit_tree(ctx, modules, fd, depth, llen, path):
    for module in modules:
        printed_header = False

        def print_header():
            bstr = ""
            b = module.search_one('belongs-to')
            if b is not None:
                bstr = " (belongs-to %s)" % b.arg
            fd.write("%s: %s%s\n" % (module.keyword, module.arg, bstr))
            printed_header = True

        chs = [ch for ch in module.i_children
               if ch.keyword in statements.data_definition_keywords]
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs if ch.arg == path[0]]
            path = path[1:]

        if len(chs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            print_children(chs, module, fd, '  ', path, 'data', depth, llen)

        mods = [module]
        for i in module.search('include'):
            subm = ctx.get_module(i.arg)
            if subm is not None:
                mods.append(subm)
        for m in mods:
            section_delimiter_printed=False
            for augment in m.search('augment'):
                if (hasattr(augment.i_target_node, 'i_module') and
                    augment.i_target_node.i_module not in modules + mods):
                    if not section_delimiter_printed:
                        fd.write('\n')
                        section_delimiter_printed = True
                    # this augment has not been printed; print it
                    if not printed_header:
                        print_header()
                        printed_header = True
                    fd.write("  augment %s:\n" % augment.arg)
                    print_children(augment.i_children, m, fd,
                                   '  ', path, 'augment', depth, llen)

        rpcs = [ch for ch in module.i_children
                if ch.keyword == 'rpc']
        if path is not None:
            if len(path) > 0:
                rpcs = [rpc for rpc in rpcs if rpc.arg == path[0]]
                path = path[1:]
            else:
                rpcs = []
        if len(rpcs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("\n  rpcs:\n")
            print_children(rpcs, module, fd, '  ', path, 'rpc', depth, llen)

        notifs = [ch for ch in module.i_children
                  if ch.keyword == 'notification']
        if path is not None:
            if len(path) > 0:
                notifs = [n for n in notifs if n.arg == path[0]]
                path = path[1:]
            else:
                notifs = []
        if len(notifs) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            fd.write("\n  notifications:\n")
            print_children(notifs, module, fd, '  ', path,
                           'notification', depth, llen)

        if ctx.opts.tree_print_groupings and len(module.i_groupings) > 0:
            if not printed_header:
                print_header()
                printed_header = True
            section_delimiter_printed = False
            for gname in module.i_groupings:
                if not section_delimiter_printed:
                    fd.write('\n')
                    section_delimiter_printed = True
                fd.write("  grouping %s\n" % gname)
                g = module.i_groupings[gname]
                print_children(g.i_children, module, fd,
                               '  ', path, 'grouping', depth, llen)

        if ctx.opts.tree_print_yang_data:
            yds = module.search(('ietf-restconf', 'yang-data'))
            if len(yds) > 0:
                if not printed_header:
                    print_header()
                    printed_header = True
                section_delimiter_printed = False
                for yd in yds:
                    if not section_delimiter_printed:
                        fd.write('\n')
                        section_delimiter_printed = True
                    fd.write("  yang-data %s:\n" % yd.arg)
                    print_children(yd.i_children, module, fd, '    ', path,
                                   'yang-data', depth, llen)

def print_children(i_children, module, fd, prefix, path, mode, depth,
                   llen, width=0):
    if depth == 0:
        if i_children: fd.write(prefix + '     ...\n')
        return
    def get_width(w, chs):
        for ch in chs:
            if ch.keyword in ['choice', 'case']:
                w = get_width(w, ch.i_children)
            else:
                if ch.i_module.i_modulename == module.i_modulename:
                    nlen = len(ch.arg)
                else:
                    nlen = len(ch.i_module.i_prefix) + 1 + len(ch.arg)
                if nlen > w:
                    w = nlen
        return w

    if width == 0:
        width = get_width(0, i_children)

    for ch in i_children:
        if ((ch.keyword == 'input' or ch.keyword == 'output') and
            len(ch.i_children) == 0):
            pass
        else:
            if (ch == i_children[-1] or
                (i_children[-1].keyword == 'output' and
                 len(i_children[-1].i_children) == 0)):
                # the last test is to detect if we print input, and the
                # next node is an empty output node; then don't add the |
                newprefix = prefix + '   '
            else:
                newprefix = prefix + '  |'
            if ch.keyword == 'input':
                mode = 'input'
            elif ch.keyword == 'output':
                mode = 'output'
            print_node(ch, module, fd, newprefix, path, mode, depth, llen,
                       width)

def print_node(s, module, fd, prefix, path, mode, depth, llen, width):
    line = "%s%s--" % (prefix[0:-1], get_status_str(s))

    brcol = len(line) + 4

    if s.i_module.i_modulename == module.i_modulename:
        name = s.arg
    else:
        name = s.i_module.i_prefix + ':' + s.arg
    flags = get_flags_str(s, mode)
    if s.keyword == 'list':
        name += '*'
        line += flags + " " + name
    elif s.keyword == 'container':
        p = s.search_one('presence')
        if p is not None:
            name += '!'
        line += flags + " " + name
    elif s.keyword  == 'choice':
        m = s.search_one('mandatory')
        if m is None or m.arg == 'false':
            line += flags + ' (' + name + ')?'
        else:
            line += flags + ' (' + name + ')'
    elif s.keyword == 'case':
        line += ':(' + name + ')'
        brcol += 1
    else:
        if s.keyword == 'leaf-list':
            name += '*'
        elif (s.keyword == 'leaf' and not hasattr(s, 'i_is_key')
              or s.keyword == 'anydata' or s.keyword == 'anyxml'):
            m = s.search_one('mandatory')
            if m is None or m.arg == 'false':
                name += '?'
        t = get_typename(s)
        if t == '':
            line += "%s %s" % (flags, name)
        elif (llen is not None and
              len(line) + len(flags) + width+1 + len(t) + 4 > llen):
            # there's no room for the type name
            if (get_leafref_path(s) is not None and
                len(t) + brcol > llen):
                # it's not even room for the leafref path; skip it
                line += "%s %-*s   leafref" % (flags, width+1, name)
            else:
                line += "%s %s" % (flags, name)
                fd.write(line + '\n')
                line = prefix + ' ' * (brcol - len(prefix)) + ' ' + t
        else:
            line += "%s %-*s   %s" % (flags, width+1, name, t)

    if s.keyword == 'list' and s.search_one('key') is not None:
        keystr = " [%s]" % re.sub('\s+', ' ', s.search_one('key').arg)
        if (llen is not None and
            len(line) + len(keystr) > llen):
            fd.write(line + '\n')
            line = prefix + ' ' * (brcol - len(prefix))
        line += keystr

    features = s.search('if-feature')
    featurenames = [f.arg for f in features]
    if hasattr(s, 'i_augment'):
        afeatures = s.i_augment.search('if-feature')
        featurenames.extend([f.arg for f in afeatures
                             if f.arg not in featurenames])

    if len(featurenames) > 0:
        fstr = " {%s}?" % ",".join(featurenames)
        if (llen is not None and len(line) + len(fstr) > llen):
            fd.write(line + '\n')
            line = prefix + ' ' * (brcol - len(prefix))
        line += fstr

    fd.write(line + '\n')
    if hasattr(s, 'i_children'):
        if depth is not None:
            depth = depth - 1
        chs = s.i_children
        if path is not None and len(path) > 0:
            chs = [ch for ch in chs
                   if ch.arg == path[0]]
            path = path[1:]
        if s.keyword in ['choice', 'case']:
            print_children(chs, module, fd, prefix, path, mode, depth,
                           llen, width)
        else:
            print_children(chs, module, fd, prefix, path, mode, depth, llen)

def get_status_str(s):
    status = s.search_one('status')
    if status is None or status.arg == 'current':
        return '+'
    elif status.arg == 'deprecated':
        return 'x'
    elif status.arg == 'obsolete':
        return 'o'

def get_flags_str(s, mode):
    if mode == 'input':
        return "-w"
    elif s.keyword in ('rpc', 'action', ('tailf-common', 'action')):
        return '-x'
    elif s.keyword == 'notification':
        return '-n'
    elif s.i_config == True:
        return 'rw'
    elif s.i_config == False or mode == 'output' or mode == 'notification':
        return 'ro'
    else:
        return '--'

def get_leafref_path(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            return t.search_one('path')
    else:
        return None

def get_typename(s):
    t = s.search_one('type')
    if t is not None:
        if t.arg == 'leafref':
            p = t.search_one('path')
            if p is not None:
                # Try to make the path as compact as possible.
                # Remove local prefixes, and only use prefix when
                # there is a module change in the path.
                target = []
                curprefix = s.i_module.i_prefix
                for name in p.arg.split('/'):
                    if name.find(":") == -1:
                        prefix = curprefix
                    else:
                        [prefix, name] = name.split(':', 1)
                    if prefix == curprefix:
                        target.append(name)
                    else:
                        target.append(prefix + ':' + name)
                        curprefix = prefix
                return "-> %s" % "/".join(target)
            else:
                return t.arg
        else:
            return t.arg
    elif s.keyword == 'anydata':
        return '<anydata>'
    elif s.keyword == 'anyxml':
        return '<anyxml>'
    else:
        return ''
<|endoftext|>