diff --git a/monkey/common/network/segmentation_utils.py b/monkey/common/network/segmentation_utils.py
index 97adbd203..6569d636b 100644
--- a/monkey/common/network/segmentation_utils.py
+++ b/monkey/common/network/segmentation_utils.py
@@ -1,14 +1,10 @@
-from common.network.network_range import NetworkRange
-
-
 def get_ip_in_src_and_not_in_dst(ip_addresses, source_subnet, target_subnet):
-    # type: (List[str], NetworkRange, NetworkRange) -> Union[str, None]
     """
     Finds an IP address in ip_addresses which is in source_subnet but not in target_subnet.
-    :param ip_addresses:    List of IP addresses to test.
-    :param source_subnet:   Subnet to want an IP to not be in.
-    :param target_subnet:   Subnet we want an IP to be in.
-    :return:    The cross segment IP if in source but not in target, else None.
+    :param ip_addresses:    List[str]: List of IP addresses to test.
+    :param source_subnet:   NetworkRange: Subnet to want an IP to not be in.
+    :param target_subnet:   NetworkRange: Subnet we want an IP to be in.
+    :return:    The cross segment IP if in source but not in target, else None. Union[str, None]
     """
     if get_ip_if_in_subnet(ip_addresses, target_subnet) is not None:
         return None