From a82da4001548b472f4df40faa174d96ec7827c3b Mon Sep 17 00:00:00 2001
From: Will Woods <wwoods@redhat.com>
Date: Wed, 27 Apr 2011 16:37:19 -0400
Subject: [PATCH] Add "getdata" method to LoraxTemplate

getdata(cmd) will return a generator that yields every token on every
line that starts with the token "cmd". getdata(cmd, mode="lines") will
yield a list for each line rather than every individual token.

this simplifies some things in __init__.py.
---
 src/pylorax/__init__.py | 24 +++++++-----------------
 src/pylorax/ltmpl.py    |  9 +++++++++
 2 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/pylorax/__init__.py b/src/pylorax/__init__.py
index d214054f..df337954 100644
--- a/src/pylorax/__init__.py
+++ b/src/pylorax/__init__.py
@@ -34,7 +34,6 @@ import os
 import ConfigParser
 import tempfile
 import shutil
-import itertools
 import glob
 import math
 import subprocess
@@ -201,24 +200,15 @@ class Lorax(BaseLoraxClass):
                   "product": self.product.name.lower() }
 
         template = ltmpl.LoraxTemplate()
-        template = template.parse(tfile, tvars)
+        template.parse(tfile, tvars)
 
-        # get required directories
         logger.info("creating tree directories")
-        dirs = [f[1:] for f in template if f[0] == "mkdir"]
-        dirs = itertools.chain.from_iterable(dirs)
-
-        # create directories
-        for d in dirs:
+        for d in template.getdata("mkdir"):
             os.makedirs(joinpaths(self.installtree.root, d))
 
-        # get list of required packages
-        logger.info("getting list of required packages")
-        required = [f[1:] for f in template if f[0] == "install"]
-        required = itertools.chain.from_iterable(required)
-
         # install packages
-        for package in required:
+        logger.info("getting list of required packages")
+        for package in template.getdata("install"):
             self.installtree.yum.install(package)
 
         skipbroken = self.conf.getboolean("yum", "skipbroken")
@@ -252,8 +242,8 @@ class Lorax(BaseLoraxClass):
         self.installtree.move_stubs()
 
         logger.info("getting list of required modules")
-        modules = [f[1:] for f in template if f[0] == "module"]
-        modules = list(itertools.chain.from_iterable(modules))
+        # Need a list to pass to cleanup_kernel_modules, not a generator
+        modules = list(template.getdata("module"))
 
         self.installtree.move_modules()
 
@@ -319,7 +309,7 @@ class Lorax(BaseLoraxClass):
         i.backup_required(self.workdir)
 
         logger.info("getting list of not required packages")
-        remove = [f[1:] for f in template if f[0] == "remove"]
+        remove = template.getdata("remove", mode="lines")
 
         rdb = {}
         order = []
diff --git a/src/pylorax/ltmpl.py b/src/pylorax/ltmpl.py
index 6485b6bd..00e870db 100644
--- a/src/pylorax/ltmpl.py
+++ b/src/pylorax/ltmpl.py
@@ -56,4 +56,13 @@ class LoraxTemplate(object):
         # split with shlex
         lines = map(shlex.split, lines)
 
+        self.lines = lines
         return lines
+
+    def getdata(self, cmd, mode="tokens"):
+        '''returns a generator that yields every token on every line starting
+        with 'cmd'. if mode is "lines", you get a list for each line instead.'''
+        if mode == "lines":
+            return (f[1:] for f in self.lines if f[0]==cmd)
+        else:
+            return (i for f in self.lines if f[0]==cmd for i in f[1:])