apply reviews from code review

airflow/utils/timezone.py

@@ -17,7 +17,7 @@
 # under the License.
 #
 import datetime as dt
-from typing import TYPE_CHECKING, Optional, Union, overload
+from typing import Optional, Union, overload
 
 import pendulum
 from dateutil.relativedelta import relativedelta
@@ -28,9 +28,6 @@
 # UTC time zone as a tzinfo instance.
 utc = pendulum.tz.timezone('UTC')
 
-if TYPE_CHECKING:
-    pass
-
 
 def is_localized(value):
     """
@@ -106,6 +103,11 @@ def make_aware(value: None, timezone: Optional[dt.tzinfo] = None) -> None:
     ...
 
 
+@overload
+def make_aware(value: DateTime, timezone: Optional[dt.tzinfo] = None) -> DateTime:
+    ...
+
+
 @overload
 def make_aware(value: dt.datetime, timezone: Optional[dt.tzinfo] = None) -> dt.datetime:
     ...
@@ -199,11 +201,16 @@ def coerce_datetime(v: None) -> None:
 
 
 @overload
-def coerce_datetime(v: dt.datetime) -> dt.datetime:
+def coerce_datetime(v: DateTime) -> DateTime:
+    ...
+
+
+@overload
+def coerce_datetime(v: dt.datetime) -> DateTime:
     ...
 
 
-def coerce_datetime(v: Optional[dt.datetime]) -> Optional[dt.datetime]:
+def coerce_datetime(v: Optional[dt.datetime]) -> Optional[DateTime]:
     """Convert whatever is passed in to an timezone-aware ``pendulum.DateTime``."""
     if v is None:
         return None