#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 . import enums, methods, types
from .__meta__ import __api_version__, __version__
from .client import session
from .client.bot import Bot
from .dispatcher.dispatcher import Dispatcher
@ -36,6 +37,3 @@ __all__ = (
"md",
"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
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 ...exceptions import TelegramNetworkError
@ -121,7 +124,12 @@ class AiohttpSession(BaseSession):
await self.close()
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
return self._session
@ -163,11 +171,21 @@ class AiohttpSession(BaseSession):
return cast(TelegramType, response.result)
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]:
if headers is None:
headers = {}
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):
yield chunk

View file

@ -158,7 +158,12 @@ class BaseSession(abc.ABC):
@abc.abstractmethod
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
"""
Stream reader

View file

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

View file

@ -49,7 +49,7 @@ dependencies = [
dynamic = ["version"]
[tool.hatch.version]
path = "aiogram/__init__.py"
path = "aiogram/__meta__.py"
[project.optional-dependencies]
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:
path = Path.cwd() / "aiogram" / "__init__.py"
path = Path.cwd() / "aiogram" / "__meta__.py"
content = path.read_text()
content = replace_line(content, API_VERSION, api_version)

View file

@ -1,5 +1,5 @@
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.client.session.base import BaseSession
@ -42,9 +42,10 @@ class MockedSession(BaseSession):
async def stream_content(
self,
url: str,
timeout: int,
chunk_size: int,
raise_for_status: bool,
headers: Optional[Dict[str, Any]] = None,
timeout: int = 30,
chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: # pragma: no cover
yield b""

View file

@ -1,6 +1,6 @@
import datetime
import json
from typing import Any, AsyncContextManager, AsyncGenerator, Optional
from typing import Any, AsyncContextManager, AsyncGenerator, Dict, Optional
from unittest.mock import AsyncMock, patch
import pytest
@ -44,7 +44,12 @@ class CustomSession(BaseSession):
assert isinstance(method, TelegramMethod)
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
assert isinstance(url, str)
assert isinstance(timeout, int)
@ -215,6 +220,7 @@ class TestBaseSession:
session = CustomSession()
stream = session.stream_content(
"https://www.python.org/static/img/python-logo.png",
headers={},
timeout=5,
chunk_size=65536,
raise_for_status=True,