diff --git a/README.md b/README.md
index 0c87d5e..4de907a 100644
--- a/README.md
+++ b/README.md
@@ -14,12 +14,16 @@ composer require --dev mathiasverraes/uptodocs
 CLI:
 
 ```
-vendor/bin/uptodocs run <markdownFile> [<preludeFile>]
+vendor/bin/uptodocs run [options] [--] <markdownFile>
 
 Arguments:
-  markdownFile  Markdown file to run.
-  preludeFile   A PHP file to run before each code block. 
-                Useful for imports and other setup code.
+  markdownFile          Markdown file to run.
+
+Options:
+  -b, --before=BEFORE   A PHP file to run before each code block. 
+                        Useful for imports and other setup code.
+  -a, --after=AFTER     A PHP file to run after each code block. 
+                        Useful for cleanup, and for running assertions.
 ```
 
 In your code: 
@@ -27,26 +31,28 @@ In your code:
 ```php
 <?php
 $upToDocs = new Verraes\UpToDocs\UpToDocs();
-$result = $upToDocs->run("README.md", "prelude.php"); // bool
+$result = $upToDocs->run("README.md"); // bool
 ```
 
 ## Example
 
-You can try it on this README file you are reading. 
-
-Run `./uptodocs run README.md` and see an error message like this: 
+You can try it on the Markdown file in the sample directory:
 
 ```
-The following code block in /Users/mathias/workspace/php/uptodocs/README.md:27 failed.
+./uptodocs run sample/docs.md --before sample/before.php
+```
+                                                
+Output:
+ 
+```
+The following code block in /Users/mathias/workspace/php/uptodocs/sample/docs.md:16 failed.
 <?php
-$upToDocs = new Verraes\UpToDocs\UpToDocs();
-$result = $upToDocs->run("README.md", "prelude.php"); // bool
+$v = multiplyy(10,2);
 
-PHP Fatal error:  Uncaught Error: Class 'Verraes\UpToDocs\UpToDocs' not found in Standard input code:4
+PHP Fatal error:  Uncaught Error: Call to undefined function multiplyy() in Standard input code:11
 Stack trace:
 #0 {main}
-  thrown in Standard input code on line 4
+  thrown in Standard input code on line 11
 ```
 
-The problem here was that `vendor/autoload.php` wasn't included, but we can fix that by adding the prelude: `./uptodocs run README.md prelude.php`. (But don't actually run that, you'll create an infinite loop!)
-
+UpToDocs discovered a typo in our sample code. Oops!
\ No newline at end of file
diff --git a/sample/before.php b/sample/before.php
new file mode 100644
index 0000000..8669fa9
--- /dev/null
+++ b/sample/before.php
@@ -0,0 +1,8 @@
+<?php declare(strict_types=1);
+
+// Typically, you'd want to do some autoloading and setup here.
+
+function multiply($x, $y)
+{
+    return $x * $y;
+}
diff --git a/sample/docs.md b/sample/docs.md
new file mode 100644
index 0000000..a9985d0
--- /dev/null
+++ b/sample/docs.md
@@ -0,0 +1,21 @@
+# UpToDocs
+
+Try UpToDocs on this Markdown file.
+
+## Examples
+
+This is how you do multiplication:
+
+```php
+<?php
+$v = multiply(1,2);
+```
+
+And another:
+
+```php
+<?php
+$v = multiplyy(10,2);
+```
+
+This last example has a typo, so UpToDocs will catch this. 
\ No newline at end of file
diff --git a/src/RunCommand.php b/src/RunCommand.php
index a92b18c..d821c01 100644
--- a/src/RunCommand.php
+++ b/src/RunCommand.php
@@ -5,6 +5,7 @@
 use Symfony\Component\Console\Command\Command;
 use Symfony\Component\Console\Input\InputArgument;
 use Symfony\Component\Console\Input\InputInterface;
+use Symfony\Component\Console\Input\InputOption;
 use Symfony\Component\Console\Output\OutputInterface;
 
 final class RunCommand extends Command
@@ -16,14 +17,20 @@ function __construct()
             ->setName('run')
             ->setDescription('Run each PHP block in a markdown file and return an error when one fails.')
             ->addArgument('markdownFile', InputArgument::REQUIRED, 'Markdown file to run.')
-            ->addArgument('preludeFile', InputArgument::OPTIONAL, 'A PHP file to run before each code block. Useful for imports and other setup code.');
+            ->addOption('before', 'b',  InputOption::VALUE_REQUIRED, 'A PHP file to run before each code block. Useful for imports and other setup code.', null)
+            ->addOption('after', 'a',  InputOption::VALUE_REQUIRED, 'A PHP file to run after each code block. Useful for cleanup, and for running assertions.', null);
     }
 
     protected function execute(InputInterface $input, OutputInterface $output)
     {
+        $beforeFile = $input->getOption('before');
+        $afterFile = $input->getOption('after');
         $markdown = $input->getArgument('markdownFile');
-        $prelude = $input->getArgument('preludeFile');
+
         $upToDocs = new UpToDocs();
-        return $upToDocs->run($markdown, $prelude) ? Command::SUCCESS : Command::FAILURE;
+        if($beforeFile) $upToDocs->before($beforeFile);
+        if($afterFile) $upToDocs->before($afterFile);
+
+        return $upToDocs->run($markdown) ? Command::SUCCESS : Command::FAILURE;
     }
 }
\ No newline at end of file
diff --git a/src/UpToDocs.php b/src/UpToDocs.php
index 39b1a20..55d8ddf 100644
--- a/src/UpToDocs.php
+++ b/src/UpToDocs.php
@@ -6,6 +6,7 @@
 use League\CommonMark\DocParser;
 use League\CommonMark\Environment;
 use League\CommonMark\Extension\CommonMarkCoreExtension;
+use League\CommonMark\Node\NodeWalkerEvent;
 use Symfony\Component\Process\Exception\ProcessFailedException;
 use Symfony\Component\Process\Process;
 
@@ -13,6 +14,9 @@ final class UpToDocs
 {
     const TIMEOUT = 5;
     private DocParser $parser;
+    private string $before = "<?php";
+    private string $after = "";
+    private ?string $workingDir = null;
 
     function __construct()
     {
@@ -22,56 +26,83 @@ function __construct()
     }
 
     /**
-     * Runs the
+     * Execute before each code block
      */
-    public function run(string $markdownFile, string $preludeFile = null): bool
+    public function before(string $beforeFile): UpToDocs
     {
-        $input = file_get_contents($markdownFile);
+        $this->before = file_get_contents($beforeFile);
+        if (is_null($this->workingDir)) {
+            $this->workingDir(dirname($beforeFile));
+        }
+        return $this;
+    }
+
+    /**
+     * The working directory to execute the code in
+     */
+    public function workingDir(string $workingDir): UpToDocs
+    {
+        $this->workingDir = $workingDir;
+        return $this;
+    }
+
+    /**
+     * Execute after each code block
+     */
+    public function after(string $afterFile): UpToDocs
+    {
+        $this->after = self::dropOpeningTag(file_get_contents($afterFile));
+        return $this;
+    }
+
+    private static function dropOpeningTag(string $phpCode): string
+    {
+        return str_replace(["<?php", "declare(strict_types=1);"], "", $phpCode);
+    }
 
-        if (is_null($preludeFile)) {
-            $prelude = "<?php\n";
-            $workingDir = dirname($markdownFile);
-        } else {
-            $prelude = file_get_contents($preludeFile);
-            $workingDir = dirname($preludeFile);
+    /**
+     * Run each code block in the markdownFile, output to STDOUT.
+     *
+     * @return bool True for successfully running all code blocks, false when one of them fails.
+     */
+    public function run(string $markdownFile): bool
+    {
+        $input = file_get_contents($markdownFile);
+        if (is_null($this->workingDir)) {
+            $this->workingDir = dirname($markdownFile);
         }
 
         $document = $this->parser->parse($input);
         $walker = $document->walker();
         while ($event = $walker->next()) {
-            $node = $event->getNode();
-            if (
-                $node instanceof FencedCode
-                && $event->isEntering()
-                && $node->getInfo() === 'php') {
+            if (self::isPHPBlock($event)) {
+                $node = $event->getNode();
 
-                $code = $prelude . "\n" . self::dropOpeningTag($node->getStringContent());
-                $process = new Process(['php'], $workingDir, null, $code, self::TIMEOUT);
+                $codeBlock = $node->getStringContent();
+                $code = $this->before . "\n" . self::dropOpeningTag($codeBlock) . "\n" . $this->after;
+                $process = new Process(['php'], $this->workingDir, null, $code, self::TIMEOUT);
 
                 try {
                     $process->mustRun();
                     echo ".";
                 } catch (ProcessFailedException $exception) {
-                    $location = realpath($markdownFile).":".$node->getStartLine();
+                    $location = realpath($markdownFile) . ":" . $node->getStartLine();
                     echo "\nThe following code block in $location failed.\n";
-                    if (!is_null($preludeFile)) {
-                        echo "(using prelude $preludeFile)\n";
-                    }
-                    echo $node->getStringContent();
+                    echo $codeBlock;
                     echo "\n";
                     echo $exception->getProcess()->getErrorOutput();
                     return false;
                 }
             }
         }
-        echo "\nOk.";
+        echo "\nOk.\n";
         return true;
     }
 
-    private static function dropOpeningTag(string $phpCode): string
+    private static function isPHPBlock(NodeWalkerEvent $event): bool
     {
-        return str_replace("<?php", "", $phpCode);
+        return $event->getNode() instanceof FencedCode
+            && $event->isEntering()
+            && $event->getNode()->getInfo() === 'php';
     }
-
-}
-
+}
\ No newline at end of file