import gzip
import sys

from . import zone_iterator, zone_dict_to_str, ZONE_FMT_STR

try:
    from colorama import Fore, init as colorama_init
except ImportError:
    HAS_COLOR = False
else:
    HAS_COLOR = True


def main():
    if HAS_COLOR:
        colors = [Fore.GREEN, Fore.MAGENTA, Fore.BLUE, Fore.CYAN, Fore.YELLOW]

        colorama_init(autoreset=True)
        unpacked_fmt = ZONE_FMT_STR.split()
        color_format = " ".join([color for segment in
                                 zip(colors, unpacked_fmt)
                                 for color in segment])

    zone_file = sys.argv[1]
    if zone_file[-2:] == 'gz':
        our_open = gzip.open
    else:
        our_open = open

    with our_open(zone_file, mode='rt') as zonefh:
        for record in zone_iterator(zonefh):
            print(zone_dict_to_str(record, fmt_str=color_format))

if __name__ == "__main__":
    main()

import argparse
import gzip
import sys

from . import zone_iterator, zone_dict_to_str, ZONE_FMT_STR

try:
    from colorama import Fore, init as colorama_init
except ImportError:
    HAS_COLOR = False
else:
    HAS_COLOR = True


def maybe_compressed_file(filename):
    if filename[-2:] == 'gz':
        our_open = gzip.open
    else:
        our_open = open

    return our_open(filename, mode='rt')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('inputs', nargs='*', type=maybe_compressed_file)
    args = parser.parse_args()

    if HAS_COLOR:
        colors = [Fore.GREEN, Fore.MAGENTA, Fore.BLUE, Fore.CYAN, Fore.YELLOW]

        colorama_init(autoreset=True)
        unpacked_fmt = ZONE_FMT_STR.split()
        color_format = " ".join([color for segment in
                                 zip(colors, unpacked_fmt)
                                 for color in segment])

    for inputfile in args.inputs:
        with inputfile as zonefh:
            for record in zone_iterator(zonefh):
                print(zone_dict_to_str(record, fmt_str=color_format))


if __name__ == "__main__":
    main()

Switch to using argparse for the main script.
Switch to using argparse for the main script.

+ import argparse
  import gzip
  import sys
  
  from . import zone_iterator, zone_dict_to_str, ZONE_FMT_STR
  
  try:
      from colorama import Fore, init as colorama_init
  except ImportError:
      HAS_COLOR = False
  else:
      HAS_COLOR = True
  
  
+ def maybe_compressed_file(filename):
+     if filename[-2:] == 'gz':
+         our_open = gzip.open
+     else:
+         our_open = open
+ 
+     return our_open(filename, mode='rt')
+ 
+ 
  def main():
+     parser = argparse.ArgumentParser()
+     parser.add_argument('inputs', nargs='*', type=maybe_compressed_file)
+     args = parser.parse_args()
+ 
      if HAS_COLOR:
          colors = [Fore.GREEN, Fore.MAGENTA, Fore.BLUE, Fore.CYAN, Fore.YELLOW]
  
          colorama_init(autoreset=True)
          unpacked_fmt = ZONE_FMT_STR.split()
          color_format = " ".join([color for segment in
                                   zip(colors, unpacked_fmt)
                                   for color in segment])
  
+     for inputfile in args.inputs:
+         with inputfile as zonefh:
+             for record in zone_iterator(zonefh):
+                 print(zone_dict_to_str(record, fmt_str=color_format))
-     zone_file = sys.argv[1]
-     if zone_file[-2:] == 'gz':
-         our_open = gzip.open
-     else:
-         our_open = open
  
-     with our_open(zone_file, mode='rt') as zonefh:
-         for record in zone_iterator(zonefh):
-             print(zone_dict_to_str(record, fmt_str=color_format))
  
  if __name__ == "__main__":
      main()
  
Switch to using argparse for the main script.
## Code Before:
import gzip
import sys

from . import zone_iterator, zone_dict_to_str, ZONE_FMT_STR

try:
    from colorama import Fore, init as colorama_init
except ImportError:
    HAS_COLOR = False
else:
    HAS_COLOR = True


def main():
    if HAS_COLOR:
        colors = [Fore.GREEN, Fore.MAGENTA, Fore.BLUE, Fore.CYAN, Fore.YELLOW]

        colorama_init(autoreset=True)
        unpacked_fmt = ZONE_FMT_STR.split()
        color_format = " ".join([color for segment in
                                 zip(colors, unpacked_fmt)
                                 for color in segment])

    zone_file = sys.argv[1]
    if zone_file[-2:] == 'gz':
        our_open = gzip.open
    else:
        our_open = open

    with our_open(zone_file, mode='rt') as zonefh:
        for record in zone_iterator(zonefh):
            print(zone_dict_to_str(record, fmt_str=color_format))

if __name__ == "__main__":
    main()

## Instruction:
Switch to using argparse for the main script.
## Code After:
import argparse
import gzip
import sys

from . import zone_iterator, zone_dict_to_str, ZONE_FMT_STR

try:
    from colorama import Fore, init as colorama_init
except ImportError:
    HAS_COLOR = False
else:
    HAS_COLOR = True


def maybe_compressed_file(filename):
    if filename[-2:] == 'gz':
        our_open = gzip.open
    else:
        our_open = open

    return our_open(filename, mode='rt')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('inputs', nargs='*', type=maybe_compressed_file)
    args = parser.parse_args()

    if HAS_COLOR:
        colors = [Fore.GREEN, Fore.MAGENTA, Fore.BLUE, Fore.CYAN, Fore.YELLOW]

        colorama_init(autoreset=True)
        unpacked_fmt = ZONE_FMT_STR.split()
        color_format = " ".join([color for segment in
                                 zip(colors, unpacked_fmt)
                                 for color in segment])

    for inputfile in args.inputs:
        with inputfile as zonefh:
            for record in zone_iterator(zonefh):
                print(zone_dict_to_str(record, fmt_str=color_format))


if __name__ == "__main__":
    main()
