Unable to pull xcom value - airflow

Question

I have a very simple DAG which includes a very simple task (PythonOperator) which gets some trivial json data from SWAPI API and returns an int. In this case the value of the int is 202.
I'm fairly sure that this int is being correctly pushed as an XCOM value because when I run the dag, select that task and look at the logs and the XCOM in the UI I see:
Furthermore when I add the line:
ti.xcom_push(key = 'height', value = height)
Into the python function which is getting the value from the API, I'm then able to see from the XCOM view for that task that a key of 'height' is being added and the value is indeed 202.
The problem is that for love nor money I cant pull that value out again and use it in another task. For example the task which needs to use that value is a PythonOperator whose function looks like:
def check_height(ti):
height = ti.xcom_pull(key = 'height', task_ids=['get_data_darth_vader'])
print(f"Height is: {height}")
I've also tried it with no key and setting the key to 'return_value', but nothing work the value is None:
[2021-09-30 21:00:35,044] {logging_mixin.py:109} INFO - Height is: [None]
[2021-09-30 21:00:35,047] {python.py:151} INFO - Done. Returned value was: None
I must be doing something wrong, but I cannot see what. I've watched a number of tutorials and read several blog posts on the subject, and can't see where what I am doing differs from the working examples.
Help!
UPDATE: Here is the whole dag
from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from datetime import datetime
import json
import requests
def get_darth_vader_height(ti):
"""
Get Darth Vader info from SWAPI
"""
response=requests.get('https://swapi.dev/api/people/4')
data=json.loads(response.text)
height=data['height']
print(f"DEBUG: {height}")
ti.xcom_push(key="height", value=height)
return height
def check_height(ti):
height=ti.xcom_pull(task_ids='task_one', key="height")
print(f"Height is: {height}")
print(str(height))
with DAG(
'my_dag',
start_date = datetime(2021,1,1),
schedule_interval="#daily",
catchup=False,
) as dag:
get_darth_vader_height = PythonOperator(
task_id='task_one',
python_callable=get_darth_vader_height
)
check_darth_vader_height = PythonOperator(
task_id='task_two',
python_callable=check_height
)
is_tall = BashOperator(
task_id='task_three',
bash_command="echo 'is tall!'"
)
is_short = BashOperator(
task_id='task_four',
bash_command="echo 'is short!'"
)
UPDATE WITH WORKING VERSION:
from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.operators.bash import BashOperator
from datetime import datetime
import json
import requests
def get_darth_vader_height(ti):
"""
Get Darth Vader info from SWAPI
"""
response=requests.get('https://swapi.dev/api/people/4')
data=json.loads(response.text)
height=data['height']
print(f"DEBUG: {height}")
ti.xcom_push(key="height", value=height)
return height
def check_height(ti):
height=ti.xcom_pull(task_ids='task_one', key="height")
print(f"Height is: {height}")
if int(height) > 200:
print('height is greater than 200')
return 'is_tall'
print('height is less than 200')
return 'is_short'
with DAG(
'my_dag',
start_date = datetime(2021,1,1),
schedule_interval="#daily",
catchup=False,
) as dag:
get_darth_vader_height = PythonOperator(
task_id='task_one',
python_callable=get_darth_vader_height
)
check_darth_vader_height = BranchPythonOperator(
task_id='task_two',
python_callable=check_height
)
is_tall = BashOperator(
task_id='task_three',
bash_command="echo 'is tall!'"
)
is_short = BashOperator(
task_id='task_four',
bash_command="echo 'is short!'"
)
get_darth_vader_height >> check_darth_vader_height
Adding the chain between the tasks fixed this issue.