diff --git a/format.sh b/format.sh
index 8df98ce..de5a6b7 100755
--- a/format.sh
+++ b/format.sh
@@ -1,15 +1,19 @@
-#!/bin/sh
+#!/bin/bash
 
-yapf \
-    --recursive \
-    --in-place \
-    --parallel \
-    .
-autoflake \
-    --recursive \
-    --in-place \
-    --remove-unused-variables \
-    --remove-all-unused-imports \
-    --expand-star-imports \
-    --remove-duplicate-keys \
-    .
+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[@]}" .
+}
+
+
+if [[ "$1" == "--check" ]]; then
+    yapf_args+=('--diff')
+    [[ "$(format | tee /dev/stderr | wc -c)" == "0" ]]
+else
+    yapf_args+=('--in-place')
+    autoflake_args+=('--in-place')
+    format
+fi