问题
A really quick quesiton, basically my program takes an inputfile which is a csv file and then converts it into a xml file. However, the name of the xml file can either be set by the user input or if the user doesn't specify the name, then the xml file will have the same name as the csv file except with a .xml extension. I need help with the latter. So far my program works when the output file is given a name but i don't know how to set the xml file name when the user doesn't input a name. I know that if the user doesn't set a name then argparse.outputfile is none and i can write a simple if statement to check but i do not know how to set the name to be the same as the csv file name after. Here is my code:
import sys, argparse
import csv
import indent
from xml.etree.ElementTree import ElementTree, Element, SubElement, Comment, tostring
parser=argparse.ArgumentParser(description='Convert wordlist text files to various formats.', prog='Text Converter')
parser.add_argument('-v','--verbose',action='store_true',dest='verbose',help='Increases messages being printed to stdout')
parser.add_argument('-c','--csv',action='store_true',dest='readcsv',help='Reads CSV file and converts to XML file with same name')
parser.add_argument('-x','--xml',action='store_true',dest='toxml',help='Convert CSV to XML with different name')
parser.add_argument('-i','--inputfile',type=argparse.FileType('r'),dest='inputfile',help='Name of file to be imported',required=True)
parser.add_argument('-o','--outputfile',type=argparse.FileType('w'),dest='outputfile',help='Output file name')
args = parser.parse_args()
def main(argv):
reader = read_csv(args.inputfile)
if args.verbose:
print ('Verbose Selected')
if args.toxml:
if args.verbose:
print ('Convert to XML Selected')
generate_xml(reader, args.outputfile)
if args.readcsv:
if args.verbose:
print ('Reading CSV file')
if not (args.toxml or args.readcsv):
parser.error('No action requested')
return 1
def read_csv(inputfile):
return list(csv.reader(inputfile))
def generate_xml(reader,outfile):
root = Element('Solution')
root.set('version','1.0')
tree = ElementTree(root)
head = SubElement(root, 'DrillHoles')
head.set('total_holes', '238')
description = SubElement(head,'description')
current_group = None
i = 0
for row in reader:
if i > 0:
x1,y1,z1,x2,y2,z2,cost = row
if current_group is None or i != current_group.text:
current_group = SubElement(description, 'hole',{'hole_id':"%s"%i})
collar = SubElement (current_group, 'collar',{'':', '.join((x1,y1,z1))}),
toe = SubElement (current_group, 'toe',{'':', '.join((x2,y2,z2))})
cost = SubElement(current_group, 'cost',{'':cost})
i+=1
indent.indent(root)
tree.write(outfile)
if (__name__ == "__main__"):
sys.exit(main(sys.argv))
回答1:
Add "import os" to the list of imports at the top of the file. Then, right after parsing, you can check and set the argument:
if args.outputfile is None:
args.outputfile = os.path.splitext(args.inputfile)[0] + '.xml'
By the way, arguments default to their long option name. You only need the 'dest' keyword when you want to use a different name. So, for example, 'verbose' could be:
parser.add_argument('-v', '--verbose', action='store_true',
help='Increases messages being printed to stdout')
EDIT: Here is the example reworked with outputfile handling and with chepner's suggestion to use positional arguments for the file names.
import os
import sys
import argparse
import csv
import indent
from xml.etree.ElementTree import ElementTree, Element, SubElement, Comment, tostring
def get_args(args):
parser=argparse.ArgumentParser(description='Convert wordlist text files to various formats.', prog='Text Converter')
parser.add_argument('-v','--verbose',action='store_true',dest='verbose',help='Increases messages being printed to stdout')
parser.add_argument('-c','--csv',action='store_true',dest='readcsv',help='Reads CSV file and converts to XML file with same name')
parser.add_argument('-x','--xml',action='store_true',dest='toxml',help='Convert CSV to XML with different name')
#parser.add_argument('-i','--inputfile',type=str,help='Name of file to be imported',required=True)
#parser.add_argument('-o','--outputfile',help='Output file name')
parser.add_argument('inputfile',type=str,help='Name of file to be imported')
parser.add_argument('outputfile',help='(optional) Output file name',nargs='?')
args = parser.parse_args()
if not (args.toxml or args.readcsv):
parser.error('No action requested')
return None
if args.outputfile is None:
args.outputfile = os.path.splitext(args.inputfile)[0] + '.xml'
return args
def main(argv):
args = get_args(argv[1:])
if args is None:
return 1
inputfile = open(args.inputfile, 'r')
outputfile = open(args.outputfile, 'w')
reader = read_csv(inputfile)
if args.verbose:
print ('Verbose Selected')
if args.toxml:
if args.verbose:
print ('Convert to XML Selected')
generate_xml(reader, outputfile)
if args.readcsv:
if args.verbose:
print ('Reading CSV file')
return 1 # you probably want to return 0 on success
def read_csv(inputfile):
return list(csv.reader(inputfile))
def generate_xml(reader,outfile):
root = Element('Solution')
root.set('version','1.0')
tree = ElementTree(root)
head = SubElement(root, 'DrillHoles')
head.set('total_holes', '238')
description = SubElement(head,'description')
current_group = None
i = 0
for row in reader:
if i > 0:
x1,y1,z1,x2,y2,z2,cost = row
if current_group is None or i != current_group.text:
current_group = SubElement(description, 'hole',{'hole_id':"%s"%i})
collar = SubElement (current_group, 'collar',{'':', '.join((x1,y1,z1))}),
toe = SubElement (current_group, 'toe',{'':', '.join((x2,y2,z2))})
cost = SubElement(current_group, 'cost',{'':cost})
i+=1
indent.indent(root)
tree.write(outfile)
if (__name__ == "__main__"):
sys.exit(main(sys.argv))
来源：https://stackoverflow.com/questions/17068536/using-argparse-to-convert-csv-to-xml-in-python