From 1621b494f49cc3e104061ab3819b5b9490b5a81f Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Wed, 21 Sep 2022 17:55:41 +0000
Subject: [PATCH] Agent: Fix mypy issues in capture_output.py

---
 monkey/infection_monkey/utils/capture_output.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/monkey/infection_monkey/utils/capture_output.py b/monkey/infection_monkey/utils/capture_output.py
index 898bd6a7e..f509fa908 100644
--- a/monkey/infection_monkey/utils/capture_output.py
+++ b/monkey/infection_monkey/utils/capture_output.py
@@ -1,9 +1,11 @@
+from __future__ import annotations
+
 import io
 import sys
 
 
 class StdoutCapture:
-    def __enter__(self) -> None:
+    def __enter__(self) -> StdoutCapture:
         self._orig_stdout = sys.stdout
         self._new_stdout = io.StringIO()
         sys.stdout = self._new_stdout