#1191 Added possibility to pass custom headers to URLInputFile object (#1197)

This commit is contained in:
Alex Root Junior 2023-06-25 01:39:26 +03:00 committed by GitHub
parent 484a61bdc1
commit d29b18da8c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 53 additions and 16 deletions

1
CHANGES/1191.feature.rst Normal file
View file

@ -0,0 +1 @@
Added possibility to pass custom headers to URLInputFile object

View file

@ -3,6 +3,7 @@ from contextlib import suppress
from aiogram.dispatcher.flags import FlagGenerator from aiogram.dispatcher.flags import FlagGenerator
from . import enums, methods, types from . import enums, methods, types
from .__meta__ import __api_version__, __version__
from .client import session from .client import session
from .client.bot import Bot from .client.bot import Bot
from .dispatcher.dispatcher import Dispatcher from .dispatcher.dispatcher import Dispatcher
@ -36,6 +37,3 @@ __all__ = (
"md", "md",
"flags", "flags",
) )
__version__ = "3.0.0b8"
__api_version__ = "6.6"

2
aiogram/__meta__.py Normal file
View file

@ -0,0 +1,2 @@
__version__ = "3.0.0b8"
__api_version__ = "6.6"

View file

@ -18,7 +18,10 @@ from typing import (
import certifi import certifi
from aiohttp import BasicAuth, ClientError, ClientSession, FormData, TCPConnector from aiohttp import BasicAuth, ClientError, ClientSession, FormData, TCPConnector
from aiohttp.hdrs import USER_AGENT
from aiohttp.http import SERVER_SOFTWARE
from aiogram.__meta__ import __version__
from aiogram.methods import TelegramMethod from aiogram.methods import TelegramMethod
from ...exceptions import TelegramNetworkError from ...exceptions import TelegramNetworkError
@ -121,7 +124,12 @@ class AiohttpSession(BaseSession):
await self.close() await self.close()
if self._session is None or self._session.closed: if self._session is None or self._session.closed:
self._session = ClientSession(connector=self._connector_type(**self._connector_init)) self._session = ClientSession(
connector=self._connector_type(**self._connector_init),
headers={
USER_AGENT: f"{SERVER_SOFTWARE} aiogram/{__version__}",
},
)
self._should_reset_connector = False self._should_reset_connector = False
return self._session return self._session
@ -163,11 +171,21 @@ class AiohttpSession(BaseSession):
return cast(TelegramType, response.result) return cast(TelegramType, response.result)
async def stream_content( async def stream_content(
self, url: str, timeout: int, chunk_size: int, raise_for_status: bool self,
url: str,
headers: Optional[Dict[str, Any]] = None,
timeout: int = 30,
chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: ) -> AsyncGenerator[bytes, None]:
if headers is None:
headers = {}
session = await self.create_session() session = await self.create_session()
async with session.get(url, timeout=timeout, raise_for_status=raise_for_status) as resp: async with session.get(
url, timeout=timeout, headers=headers, raise_for_status=raise_for_status
) as resp:
async for chunk in resp.content.iter_chunked(chunk_size): async for chunk in resp.content.iter_chunked(chunk_size):
yield chunk yield chunk

View file

@ -158,7 +158,12 @@ class BaseSession(abc.ABC):
@abc.abstractmethod @abc.abstractmethod
async def stream_content( async def stream_content(
self, url: str, timeout: int, chunk_size: int, raise_for_status: bool self,
url: str,
headers: Optional[Dict[str, Any]] = None,
timeout: int = 30,
chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: # pragma: no cover ) -> AsyncGenerator[bytes, None]: # pragma: no cover
""" """
Stream reader Stream reader

View file

@ -4,7 +4,7 @@ import io
import os import os
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from pathlib import Path from pathlib import Path
from typing import AsyncGenerator, AsyncIterator, Iterator, Optional, Union from typing import Any, AsyncGenerator, AsyncIterator, Dict, Iterator, Optional, Union
import aiofiles import aiofiles
@ -114,6 +114,7 @@ class URLInputFile(InputFile):
def __init__( def __init__(
self, self,
url: str, url: str,
headers: Optional[Dict[str, Any]] = None,
filename: Optional[str] = None, filename: Optional[str] = None,
chunk_size: int = DEFAULT_CHUNK_SIZE, chunk_size: int = DEFAULT_CHUNK_SIZE,
timeout: int = 30, timeout: int = 30,
@ -122,12 +123,16 @@ class URLInputFile(InputFile):
Represents object for streaming files from internet Represents object for streaming files from internet
:param url: URL in internet :param url: URL in internet
:param headers: HTTP Headers
:param filename: Filename to be propagated to telegram. :param filename: Filename to be propagated to telegram.
:param chunk_size: Uploading chunk size :param chunk_size: Uploading chunk size
""" """
super().__init__(filename=filename, chunk_size=chunk_size) super().__init__(filename=filename, chunk_size=chunk_size)
if headers is None:
headers = {}
self.url = url self.url = url
self.headers = headers
self.timeout = timeout self.timeout = timeout
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]: async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
@ -136,6 +141,7 @@ class URLInputFile(InputFile):
bot = Bot.get_current(no_error=False) bot = Bot.get_current(no_error=False)
stream = bot.session.stream_content( stream = bot.session.stream_content(
url=self.url, url=self.url,
headers=self.headers,
timeout=self.timeout, timeout=self.timeout,
chunk_size=self.chunk_size, chunk_size=self.chunk_size,
raise_for_status=True, raise_for_status=True,

View file

@ -49,7 +49,7 @@ dependencies = [
dynamic = ["version"] dynamic = ["version"]
[tool.hatch.version] [tool.hatch.version]
path = "aiogram/__init__.py" path = "aiogram/__meta__.py"
[project.optional-dependencies] [project.optional-dependencies]
fast = [ fast = [

View file

@ -46,7 +46,7 @@ def replace_line(content: str, pattern: re.Pattern, new_value: str) -> str:
def write_package_meta(api_version: str) -> None: def write_package_meta(api_version: str) -> None:
path = Path.cwd() / "aiogram" / "__init__.py" path = Path.cwd() / "aiogram" / "__meta__.py"
content = path.read_text() content = path.read_text()
content = replace_line(content, API_VERSION, api_version) content = replace_line(content, API_VERSION, api_version)

View file

@ -1,5 +1,5 @@
from collections import deque from collections import deque
from typing import TYPE_CHECKING, AsyncGenerator, Deque, Optional, Type from typing import TYPE_CHECKING, Any, AsyncGenerator, Deque, Dict, Optional, Type
from aiogram import Bot from aiogram import Bot
from aiogram.client.session.base import BaseSession from aiogram.client.session.base import BaseSession
@ -42,9 +42,10 @@ class MockedSession(BaseSession):
async def stream_content( async def stream_content(
self, self,
url: str, url: str,
timeout: int, headers: Optional[Dict[str, Any]] = None,
chunk_size: int, timeout: int = 30,
raise_for_status: bool, chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: # pragma: no cover ) -> AsyncGenerator[bytes, None]: # pragma: no cover
yield b"" yield b""

View file

@ -1,6 +1,6 @@
import datetime import datetime
import json import json
from typing import Any, AsyncContextManager, AsyncGenerator, Optional from typing import Any, AsyncContextManager, AsyncGenerator, Dict, Optional
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
import pytest import pytest
@ -44,7 +44,12 @@ class CustomSession(BaseSession):
assert isinstance(method, TelegramMethod) assert isinstance(method, TelegramMethod)
async def stream_content( async def stream_content(
self, url: str, timeout: int, chunk_size: int, raise_for_status: bool self,
url: str,
headers: Optional[Dict[str, Any]] = None,
timeout: int = 30,
chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: # pragma: no cover ) -> AsyncGenerator[bytes, None]: # pragma: no cover
assert isinstance(url, str) assert isinstance(url, str)
assert isinstance(timeout, int) assert isinstance(timeout, int)
@ -215,6 +220,7 @@ class TestBaseSession:
session = CustomSession() session = CustomSession()
stream = session.stream_content( stream = session.stream_content(
"https://www.python.org/static/img/python-logo.png", "https://www.python.org/static/img/python-logo.png",
headers={},
timeout=5, timeout=5,
chunk_size=65536, chunk_size=65536,
raise_for_status=True, raise_for_status=True,