[GitHub] [airflow] ashb commented on a change in pull request #11486: Annotate DagRun methods with return types.

[GitBox <gi...@apache.org>]

ashb commented on a change in pull request #11486:
URL: https://github.com/apache/airflow/pull/11486#discussion_r507626299



##########
File path: airflow/models/dagrun.py
##########
@@ -310,15 +315,13 @@ def get_task_instance(self, task_id: str, session: Session = None):
         :param session: Sqlalchemy ORM Session
         :type session: Session
         """
-        ti = session.query(TI).filter(
+        return session.query(TI).filter(
             TI.dag_id == self.dag_id,
             TI.execution_date == self.execution_date,
             TI.task_id == task_id
         ).first()
 
-        return ti
-
-    def get_dag(self):
+    def get_dag(self) -> 'DAG':

Review comment:
       Since you have the `if TYPE_CHECKING` can this use the non-string form?

##########
File path: airflow/models/dagrun.py
##########
@@ -301,7 +306,7 @@ def get_task_instances(self, state=None, session=None):
         return tis.all()
 
     @provide_session
-    def get_task_instance(self, task_id: str, session: Session = None):
+    def get_task_instance(self, task_id: str, session: Session = None) -> TI:

Review comment:
       I _think_ this might be `Optional[TI]`




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org