如何在Jupyter Notebook中使用asyncio上下文管理器创建异步任务并修复数据返回异常

我有一个继承自Fred基类的Series类，该类包含多个从不同API端点获取数据并将其处理为对象属性的方法。Fred基类仅包含获取API密钥、格式化请求及向服务器提交请求的方法。为提升数据获取效率，我用asyncio编写了所有相关的获取方法，在Python脚本或IPython控制台中运行正常，但在Jupyter Notebook中使用时，虽能将任务添加到现有事件循环，异步任务却始终无法正确返回数据。

相关代码 #

核心方法：run_fetch_data #

该方法在初始化时运行，也可重新调用：

def run_fetch_data(self, *args):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:  # 'RuntimeError: There is no current event loop...'
            loop = None

        if loop and loop.is_running():
            print('Async event loop already running. Adding coroutine to the event loop.')
            loop.create_task(self.async_fetch_data(*args, current_loop=loop))
        else:
            print('Starting new event loop')
            asyncio.run(self.async_fetch_data(*args))
        return self

辅助函数1：async_fetch_data #

（注意：若未提供kwargs，可返回正确数据）

async def async_fetch_data(self, *args, **kwargs):
        if not args:
            tasks = [self._fetch_series_info,
                     self._fetch_series_release_info,
                     self._fetch_series_vintagedates,
                     self._fetch_series_categories,
                     self._fetch_series_tags,
                     self._fetch_series_observations]
        else:
            tasks = []
            if 'info' in args:
                tasks.append(self._fetch_series_info)
            else:
                pass
            if 'release_info' in args:
                tasks.append(self._fetch_series_release_info)
            else:
                pass
            if 'categories' in args:
                tasks.append(self._fetch_series_categories)
            else:
                pass
            if 'observations' in args:
                tasks.append(self._fetch_series_observations)
            else:
                pass
            if 'vintagedates' in args:
                tasks.append(self._fetch_series_vintagedates)
            if 'tags' in args:
                tasks.append(self._fetch_series_tags)
            else:
                pass
        if not kwargs or kwargs is None:
            async with aiohttp.ClientSession(self._url_base) as session:
                async with asyncio.TaskGroup() as tg:
                    for task in tasks:
                        tg.create_task(task(session))
        elif 'current_loop' in kwargs.keys():
            loop = kwargs['current_loop']
            async with aiohttp.ClientSession(self._url_base) as session:
                for task in tasks:
                    loop.create_task(task(session))
        else:
            raise ValueError('Keyword argument "current_loop" not found in kwargs')
        return self

辅助函数2：_fetch_series_info #

任务中调用的其他函数与此类似：

async def _fetch_series_info(self, session):
        try:
            raw_response = await self._async_fetch_data(session, self._series_info_path,
                                                  f'series_id={self.series_id}',
                                                        **self._url_options)
            if isinstance(raw_response, dict):
                raw_response = raw_response['seriess'][0]
            elif isinstance(raw_response, ET.Element):
                raw_response = raw_response[0].attrib
            else:
                raise TypeError('Fred response was neither JSON nor XML, series info get failed')
            setattr(self, 'series_name', raw_response['title'])
            setattr(self, 'observation_start', dt.date.fromisoformat(raw_response['observation_start']))
            setattr(self, 'observation_end', dt.date.fromisoformat(raw_response['observation_end']))
            setattr(self, 'frequency', raw_response['frequency_short'])
            setattr(self, 'units', raw_response['units_short'])
            setattr(self, 'seasonality_adjustment', raw_response['seasonal_adjustment_short'])
            setattr(self, 'series_notes', raw_response['notes'])
            setattr(self, 'update_info', dt.datetime.fromisoformat(raw_response['last_updated']))
        except TypeError:
            print(f'Error fetching info for series {self.series_id}')
        except HTTPError as http_msg:
            print(http_msg)
        finally:
            return self

辅助函数3：Fred基类中的_async_fetch_data #

async def _async_fetch_data(self, session: aiohttp.ClientSession,
                                query_path: str, query: str = '', **kwargs) -> dict | ET.ElementTree | ValueError:
        """
        Internal function used to get and parse requests from FRED server

        Parameters
        ----------
        target_url : str
            the target url including any options for the retrieval of information from FRED

        Returns
        -------
        response : dict, XT.Element
            JSON: a dictionary containing the searched data.
            XML: an Element Tree or Element Tree.Element containing searched data.
        """
        request_parse = url_parse.urlparse(Fred._url_base)
        query_options = {'api_key': self.api_key,
                         'file_type': self._default_response_type,
                         'max_response_per_request': self._max_response_per_request,
                         'realtime_start': (self.realtime_start if self.realtime_start is not None else self._earliest_realtime_start),
                         'realtime_end': (self.realtime_end if self.realtime_end is not None else self._latest_realtime_end)}
        if self.observation_end is not None:
            query_options['observation_end'] = str(self.observation_end)
        if self.observation_start is not None:
            query_options['observation_start'] = str(self.observation_start)
        for kw, arg in kwargs.items():
            if self.check_optional_request_parameters(kw, arg):
                query_options[kw] = arg
        target_path = ''.join([query_path, '?', query, '&', url_parse.urlencode(query_options)])
        target_url = url_parse.urlunparse((url_parse.urlparse(Fred._url_base)
                         ._replace(path=url_parse.urljoin(request_parse.path, query_path))
                         ._replace(query=''.join([query,'&',url_parse.urlencode(query_options)]))))
        try:
            async with session.get(target_path) as raw_response:
                if query_options['file_type'] == 'json':
                    response = await raw_response.json()
                else:
                    response = ET.fromstring(await raw_response.text())
        except HTTPError as msg:
            response = ET.fromstring(msg.read())
            raise ValueError(response.get('message'))
        return response

核心问题 #

在Jupyter Notebook环境下，现有事件循环已启动时，你调用loop.create_task()创建子任务后，没有等待这些任务完成就结束了async_fetch_data函数。此时async with aiohttp.ClientSession块会直接关闭会话，导致子任务中的HTTP请求还没完成就被中断，自然无法返回数据。

修复方案 #

1. 统一任务等待逻辑 #

删除冗余的分支判断，不管是否使用现有循环，都用TaskGroup等待所有子任务完成：

async def async_fetch_data(self, *args, **kwargs):
        if not args:
            tasks = [self._fetch_series_info,
                     self._fetch_series_release_info,
                     self._fetch_series_vintagedates,
                     self._fetch_series_categories,
                     self._fetch_series_tags,
                     self._fetch_series_observations]
        else:
            tasks = []
            if 'info' in args:
                tasks.append(self._fetch_series_info)
            if 'release_info' in args:
                tasks.append(self._fetch_series_release_info)
            if 'categories' in args:
                tasks.append(self._fetch_series_categories)
            if 'observations' in args:
                tasks.append(self._fetch_series_observations)
            if 'vintagedates' in args:
                tasks.append(self._fetch_series_vintagedates)
            if 'tags' in args:
                tasks.append(self._fetch_series_tags)
        
        async with aiohttp.ClientSession(self._url_base) as session:
            async with asyncio.TaskGroup() as tg:
                for task in tasks:
                    tg.create_task(task(session))
        
        return self

2. 调整run_fetch_data的返回逻辑 #

在Jupyter中，同步函数里创建异步任务后不会自动等待完成，所以可以返回任务对象让用户手动等待：

def run_fetch_data(self, *args):
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop and loop.is_running():
            print('Async event loop already running. Scheduling coroutine.')
            task = loop.create_task(self.async_fetch_data(*args))
            return task
        else:
            print('Starting new event loop')
            asyncio.run(self.async_fetch_data(*args))
        return self

3. Jupyter中的正确使用方式 #

在Jupyter Notebook里，推荐直接await异步方法，确保任务完成后再继续：

# 方式1：直接调用异步方法并等待
series = Series(...)
await series.async_fetch_data()

# 方式2：使用run_fetch_data返回的任务对象
task = series.run_fetch_data()
await task

额外优化 #

• 移除代码中冗余的else: pass分支，简化逻辑
• 确保所有异步任务都被正确等待，避免会话提前关闭
• 在Jupyter中优先使用await调用异步方法，逻辑更直观

内容的提问来源于stack exchange，提问作者crm5024