diff --git a/format.sh b/format.sh
index de5a6b7..1514aee 100755
--- a/format.sh
+++ b/format.sh
@@ -4,16 +4,24 @@ yapf_args=('--recursive' '--parallel')
 autoflake_args=('--recursive' '--remove-unused-variables' '--remove-all-unused-imports' '--expand-star-imports' '--remove-duplicate-keys')
 
 format() {
-    yapf "${yapf_args[@]}"  .
-    autoflake "${autoflake_args[@]}" .
+    files=("$@")
+    if [[ -z "${files[@]}" ]]; then
+        files=(".")
+    fi
+
+    yapf "${yapf_args[@]}"  "${files[@]}"
+    autoflake "${autoflake_args[@]}" "${files[@]}"
 }
 
 
+
+
 if [[ "$1" == "--check" ]]; then
     yapf_args+=('--diff')
-    [[ "$(format | tee /dev/stderr | wc -c)" == "0" ]]
+    shift
+    [[ "$(format "$@" | tee /dev/stderr | wc -c)" == "0" ]]
 else
     yapf_args+=('--in-place')
     autoflake_args+=('--in-place')
-    format
+    format "$@"
 fi