如何實現標準輸出和文件寫入基於參數輸入

Question

我有一個看起來輸入文件中像this（infile.txt）：

a x 
b y 
c z 

我想要實現一個程序，允許用戶寫入STDOUT或文件取決於命令：

python mycode.py infile.txt outfile.txt 

將寫入文件。

而與此

python mycode.py infile.txt #2nd case 

將寫入標準輸出。

我堅持用這個代碼：

import sys 
import csv 

nof_args = len(sys.argv) 
infile = sys.argv[1] 

print nof_args 
outfile = '' 
if nof_args == 3: 
    outfile = sys.argv[2] 

# for some reason infile is so large 
# so we can't save it to data structure (e.g. list) for further processing 
with open(infile, 'rU') as tsvfile: 
    tabreader = csv.reader(tsvfile, delimiter=' ') 

    with open(outfile, 'w') as file: 
     for line in tabreader: 
      outline = "__".join(line) 
      # and more processing 
      if nof_args == 3: 
       file.write(outline + "\n") 
      else: 
       print outline 
    file.close() 

當使用第二種情況是產生

Traceback (most recent call last): 
    File "test.py", line 18, in <module> 
    with open(outfile, 'w') as file: 
IOError: [Errno 2] No such file or directory: '' 

什麼來實現它的更好的辦法？

Answer 1

你可以試試這個：

import sys 

if write_to_file: 
    out = open(file_name, 'w') 
else: 
    out = sys.stdout 

# or a one-liner: 
# out = open(file_name, 'w') if write_to_file else sys.stdout 

for stuff in data: 
    out.write(stuff) 

out.flush() # cannot close stdout 

# Python deals with open files automatically 

你也可以用這個來代替out.flush()：

try: 
    out.close() 
except AttributeError: 
    pass 

這看起來有點醜到我，所以，flush將是多麼好。