diff --git a/mlir/utils/generate-test-checks.py b/mlir/utils/generate-test-checks.py
index e0115e51f269..7c553135cb6f 100755
--- a/mlir/utils/generate-test-checks.py
+++ b/mlir/utils/generate-test-checks.py
@@ -35,7 +35,6 @@ import argparse
 import os  # Used to advertise this file's name ("autogenerated_note").
 import re
 import sys
-import string
 
 ADVERT = '// NOTE: Assertions have been autogenerated by '
 
@@ -114,9 +113,8 @@ def preprocess_line(line):
 
 
 def main():
-  from argparse import RawTextHelpFormatter
   parser = argparse.ArgumentParser(
-      description=__doc__, formatter_class=RawTextHelpFormatter)
+      description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
   parser.add_argument(
       '--check-prefix', default='CHECK', help='Prefix to use from check file.')
   parser.add_argument(