Fixing various issues in utils folder reported by MyPy

airflow/utils/dates.py

@@ -105,7 +105,7 @@ def date_range(
     else:
         raise Exception("Wait. delta must be either datetime.timedelta or cron expression as str")
 
-    dates = []
+    dates: List[datetime] = []
     if end_date:
         if timezone.is_naive(start_date) and not timezone.is_naive(end_date):
             end_date = timezone.make_naive(end_date, time_zone)

airflow/utils/dot_renderer.py

@@ -111,7 +111,7 @@ def _draw_task_group(
 
 
 def _draw_nodes(
-    node: DependencyMixin, parent_graph: graphviz.Digraph, states_by_task_id: Dict[str, str]
+    node: DependencyMixin, parent_graph: graphviz.Digraph, states_by_task_id: Optional[Dict[str, str]]
 ) -> None:
     """Draw the node and its children on the given parent_graph recursively."""
     if isinstance(node, BaseOperator):
@@ -136,12 +136,12 @@ def _draw_nodes(
                 _draw_task_group(node, sub, states_by_task_id)
 
 
-def render_dag_dependencies(deps: Optional[Dict[str, List['DagDependency']]]) -> graphviz.Digraph:
+def render_dag_dependencies(deps: Dict[str, List['DagDependency']]) -> graphviz.Digraph:
     """
     Renders the DAG dependency to the DOT object.
 
     :param deps: List of DAG dependencies
-    :type deps: Optional[List[DagDependency]]
+    :type deps: Dict[str, List['DagDependency']]
     :return: Graphviz object
     :rtype: graphviz.Digraph
     """

airflow/utils/file.py

@@ -83,9 +83,10 @@ def correct_maybe_zipped(fileloc: Union[None, str, Path]) -> Union[None, str, Pa
     If the path contains a folder with a .zip suffix, then
     the folder is treated as a zip archive and path to zip is returned.
     """
-    if not fileloc:
-        return fileloc
-    _, archive, _ = ZIP_REGEX.search(fileloc).groups()
+    archive = None
+    matched = ZIP_REGEX.search(str(fileloc))
+    if matched:
+        _, archive, _ = matched.groups()
     if archive and zipfile.is_zipfile(archive):
         return archive
     else:

airflow/utils/log/file_task_handler.py

@@ -25,6 +25,7 @@
 from itsdangerous import TimedJSONWebSignatureSerializer
 
 from airflow.configuration import AirflowConfigException, conf
+from airflow.exceptions import AirflowException
 from airflow.utils.context import Context
 from airflow.utils.helpers import parse_template_string, render_template_to_string
 from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler
@@ -82,13 +83,15 @@ def _render_filename(self, ti: "TaskInstance", try_number: int) -> str:
                 context = Context(ti=ti, ts=ti.get_dagrun().logical_date.isoformat())
             context["try_number"] = try_number
             return render_template_to_string(self.filename_jinja_template, context)
-
-        return self.filename_template.format(
-            dag_id=ti.dag_id,
-            task_id=ti.task_id,
-            execution_date=ti.get_dagrun().logical_date.isoformat(),
-            try_number=try_number,
-        )
+        elif self.filename_template:
+            return self.filename_template.format(
+                dag_id=ti.dag_id,
+                task_id=ti.task_id,
+                execution_date=ti.get_dagrun().logical_date.isoformat(),
+                try_number=try_number,
+            )
+        else:
+            raise AirflowException("self.filename_jinja_template or self.filename_template not defined")
 
     def _read_grouped_logs(self):
         return False

airflow/utils/python_virtualenv.py

@@ -20,7 +20,7 @@
 import os
 import sys
 from collections import deque
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import jinja2
 
@@ -117,6 +117,7 @@ def write_python_script(
     :param render_template_as_native_obj: If ``True``, rendered Jinja template would be converted
         to a native Python object
     """
+    template_env: Union[jinja2.Environment, jinja2.nativetypes.NativeEnvironment]
     template_loader = jinja2.FileSystemLoader(searchpath=os.path.dirname(__file__))
     if render_template_as_native_obj:
         template_env = jinja2.nativetypes.NativeEnvironment(

airflow/utils/timezone.py

@@ -17,7 +17,8 @@
 # under the License.
 #
 import datetime as dt
-from typing import TYPE_CHECKING, Optional, Union, overload
+from datetime import tzinfo
+from typing import TYPE_CHECKING, Optional, Union
 
 import pendulum
 from dateutil.relativedelta import relativedelta
@@ -101,17 +102,7 @@ def convert_to_utc(value):
     return value.astimezone(utc)
 
 
-@overload
-def make_aware(v: None, timezone: Optional["Timezone"] = None) -> None:
-    ...
-
-
-@overload
-def make_aware(v: dt.datetime, timezone: Optional["Timezone"] = None) -> dt.datetime:
-    ...
-
-
-def make_aware(value: Optional[dt.datetime], timezone: Optional["Timezone"] = None) -> Optional[dt.datetime]:
+def make_aware(value: dt.datetime, timezone: Optional[Union["Timezone", "tzinfo"]] = None) -> dt.datetime:
     """
     Make a naive datetime.datetime in a given time zone aware.
 
@@ -133,10 +124,10 @@ def make_aware(value: Optional[dt.datetime], timezone: Optional["Timezone"] = No
         value = value.replace(fold=1)
     if hasattr(timezone, 'localize'):
         # This method is available for pytz time zones.
-        return timezone.localize(value)
+        return getattr(timezone, "localize")(value)
     elif hasattr(timezone, 'convert'):
         # For pendulum
-        return timezone.convert(value)
+        return getattr(timezone, 'convert')(value)
     else:
         # This may be wrong around DST changes!
         return value.replace(tzinfo=timezone)
@@ -189,24 +180,14 @@ def parse(string: str, timezone=None) -> DateTime:
     return pendulum.parse(string, tz=timezone or TIMEZONE, strict=False)  # type: ignore
 
 
-@overload
-def coerce_datetime(v: None) -> None:
-    ...
-
-
-@overload
-def coerce_datetime(v: dt.datetime) -> DateTime:
-    ...
-
-
-def coerce_datetime(v: Optional[dt.datetime]) -> Optional[DateTime]:
-    """Convert whatever is passed in to an timezone-aware ``pendulum.DateTime``."""
+def coerce_datetime(v: Optional[dt.datetime]) -> Optional[dt.datetime]:
+    """Convert whatever is passed in to a timezone-aware ``pendulum.DateTime``."""
     if v is None:
         return None
     if isinstance(v, DateTime):
         return v if v.tzinfo else make_aware(v)
     # Only dt.datetime is left here
-    return pendulum.instance(v if v.tzinfo else make_aware(v))
+    return pendulum.instance(v if v.tzinfo else make_aware(v))  # type: ignore
 
 
 def td_format(td_object: Union[None, dt.timedelta, float, int]) -> Optional[str]:
@@ -220,7 +201,7 @@ def td_format(td_object: Union[None, dt.timedelta, float, int]) -> Optional[str]
     if isinstance(td_object, dt.timedelta):
         delta = relativedelta() + td_object
     else:
-        delta = relativedelta(seconds=td_object)
+        delta = relativedelta(seconds=td_object)  # type: ignore
     # relativedelta for timedelta cannot convert days to months
     # so calculate months by assuming 30 day months and normalize
     months, delta.days = divmod(delta.days, 30)

tests/utils/test_file.py

@@ -25,6 +25,15 @@
 
 
 class TestCorrectMaybeZipped(unittest.TestCase):
+    @mock.patch("zipfile.is_zipfile")
+    def test_correct_maybe_zipped_empty(self, mocked_is_zipfile):
+        path = ''
+        mocked_is_zipfile.return_value = False
+
+        dag_folder = correct_maybe_zipped(path)
+
+        assert dag_folder == path
+
     @mock.patch("zipfile.is_zipfile")
     def test_correct_maybe_zipped_normal_file(self, mocked_is_zipfile):
         path = '/path/to/some/file.txt'