Note: This was originally posted by an inactive account. Content was preserved by moving under an admin account.
Originally posted by: dkhuon@gogoair.comI got this code which was doing a 1-on-1 file zipping: A.csv -> A.zip, B.csv -> B.zip.
I like to override it to use a different filename for the zip file. I could, for instance ask it to zip A.csv -> AX.zip, B.csv -> BX.zip, C.csv -> CZ.zip. I also could use it where the target zip file is the same for all the input.
The following code was modified (in pump() and openArchive() only) to use the zip filename from the input stream. Here is the result:
What is right: It creates the zip file as expected!
What is wrong: The output has more rows than the input. I need the output to be the same as input. See example below.
Work-around solution: I could ignore the output, and re-use the same input file for sub-sequent nodes. Still like to clean-up this code though.
Please help.
Input:
A.csv T.csv
B.csv T.csv
C.csv T.csv
Output:
A.csv T.csv
A.csv T.csv
B.csv T.csv
A.csv T.csv
B.csv T.csv
C.csv T.csv
Code:
import braininfo
import os
import zipfile
import gzip
#import bz2
def setup(brainNodeControlObj, BrainNodeClass):
''' Must return a class that inherits from BrainNodeClass. '''
class BrainNode(BrainNodeClass):
def initialize(self):
''' Called at node initialization, before first pump.'''
super(BrainNode, self).initialize()
# Initialize Parameter Values
self.archiveType = self.properties.getString("ls.brain.node.CreateZip File.ArchiveType")
self.fileExpr = self.properties.getString("ls.brain.node.CreateZip File.FileExpr")
self.useCompression = self.properties.getBool("ls.brain.node.CreateZipFi le.UseCompression")
self.keepOriginal = self.properties.getBool("ls.brain.node.CreateZipFi le.KeepOriginal")
self.errorIfNoExist = self.properties.getBool("ls.brain.node.CreateZipFi le.ErrorIfNoExist")
# Initialize Input Details
self.input = self.inputs[0]
self.im = self.input.metadata
# Initialize Output Details
self.output = self.outputs[0]
self.om = self.newMetadata()
self.om.append("FileName","string")
self.om.append("ArchiveName","string")
self.om.append("DateTime","string")
self.om.append("FileSize","string")
self.output.metadata = self.om
def finalize(self, val):
''' Called at node end, after last pump call. '''
super(BrainNode, self).finalize(val)
def openArchive(self):
''' Call this to open the archive file. '''
filePath, fileName = os.path.split(self.file)
fileBaseName, fileExtension = os.path.splitext(fileName)
if self.archiveType == 'zip':
self.archiveFile = self.ArchiveName # Use different name for zip
# self.archiveFile = filePath+'/'+fileBaseName+'.'+self.archiveType
elif self.archiveType == 'gzip':
self.archiveFile = filePath+'/'+fileBaseName+fileExtension+'.gz'
elif self.archiveType == 'bz2':
self.archiveFile = filePath+'/'+fileBaseName+'.'+self.archiveType
if os.access(self.archiveFile, os.F_OK):
self.archiveAvailable = True
else:
self.archiveAvailable = False
# Open the Archive File for writing or appending
if self.archiveType == 'zip':
if self.archiveAvailable and zipfile.is_zipfile(self.archiveFile):
if self.useCompression:
self.archive = zipfile.ZipFile(self.archiveFile, "a", zipfile.ZIP_DEFLATED)
else:
self.archive = zipfile.ZipFile(self.archiveFile, "a")
else:
if self.useCompression:
self.archive = zipfile.ZipFile(self.archiveFile, "w", zipfile.ZIP_DEFLATED)
else:
self.archive = zipfile.ZipFile(self.archiveFile, "w")
elif self.archiveType == 'gzip':
if self.archiveAvailable:
self.archive = gzip.GzipFile(self.archiveFile, "ab")
else:
self.archive = gzip.GzipFile(self.archiveFile, "wb")
elif self.archiveType == 'bz2':
self.archive = bz2.BZ2File(self.archiveFile, "w")
def closeArchive(self):
''' Call this to close the archive file. '''
# Close the Archive File
self.archive.close()
# Open the Archive File to get info
# DK: still miscount the files - 2/12/2014
if self.archiveType == 'zip':
if os.access(self.archiveFile, os.F_OK) and zipfile.is_zipfile(self.archiveFile):
self.archive = zipfile.ZipFile(self.archiveFile, "r")
outRec = self.output.newRecord()
# Archive file info
for info in self.archive.infolist():
outRec["FileName"] = info.filename
outRec["ArchiveName"] = self.archiveFile
outRec["DateTime"] = info.date_time
outRec["FileSize"] = info.file_size
self.output.write(outRec)
# Check the Archive File for bad files
badFile = self.archive.testzip()
# Close the Archive File again
self.archive.close()
if badFile == None:
pass
else:
raise braininfo.BrainNodeException, 'Archive contains bad file: "' + str(badFile) + '"'
elif self.archiveType == 'gzip' or self.archiveType == 'bz2':
if os.access(self.archiveFile, os.F_OK):
outRec = self.output.newRecord()
# Archive file info
info = os.stat(self.archiveFile)
outRec["FileName"] = self.file
outRec["ArchiveName"] = self.archiveFile
outRec["DateTime"] = info.st_mtime
outRec["FileSize"] = info.st_size
self.output.write(outRec)
def pump(self, quant):
''' Will continue to be called until False is returned.
Only a small amount of work should be done in pump without checking
quant.permitsRunning. When that returns False, return from pump
with a True value to be called again, a False value when finished.'''
while quant.permitsRunning(self):
# your code here
# Read first input record
rec = self.input.read()
# If no more records, exit the pump
if rec is None:
return False
# Get Full Files Name + Path
self.file = rec[self.fileExpr]
self.ArchiveName = rec[self.properties.getString("ls.brain.node.CreateZip File.ZipName")] # store it for openArchive
# Open Archive File
self.openArchive()
# Add the file to the archive
if os.access(self.file, os.F_OK):
if self.archiveType == 'zip':
self.archive.write(self.file, os.path.basename(self.file))
elif self.archiveType == 'gzip' or self.archiveType == 'bz2':
openFile = open(self.file,'rb')
for line in openFile:
self.archive.write(line)
openFile.close()
# Delete file if asked
if not self.keepOriginal:
os.remove(self.file)
else:
if self.errorIfNoExist:
raise braininfo.BrainNodeException, 'File does not exist: "' + self.file + '"'
# Close Archive File
self.closeArchive()
# finished your quantum, you want pump called again
return True
# If you implement your logic here, you don't need code below. If you want
# to inherit, and add more logic at next level, uncomment below code
# py3File = brainNodeControlObj.properties.getString("ls.brain .node.BrainPython.python3ImplementationFile")
# braininfo.loadModule(py3File, "py3", brainNodeControlObj)
# import py3
# BrainNodeImpl = py3.setup(brainNodeControlObj, BrainNode)
# return BrainNodeImpl
return BrainNode