From 5478e01b35934e68e5e34cf14745cb3c29827f5b Mon Sep 17 00:00:00 2001 From: CRIMINAL Date: Sun, 13 Jun 2021 02:54:25 +0100 Subject: [PATCH] Less memory will be used when downloading This actually lowers memory usage when downloading files so you can run with max threads now. Fixed https://github.com/DIGITALCRIMINAL/OnlyFans/issues/1331 --- apis/api_helper.py | 95 +++++++++++++++++++++--------------------- helpers/main_helper.py | 26 ++++++++---- modules/onlyfans.py | 17 ++++---- start_ofd.py | 4 +- 4 files changed, 77 insertions(+), 65 deletions(-) diff --git a/apis/api_helper.py b/apis/api_helper.py index dba05ec22..942395dd0 100644 --- a/apis/api_helper.py +++ b/apis/api_helper.py @@ -19,7 +19,8 @@ import requests from aiohttp import ClientSession from aiohttp.client_exceptions import (ClientConnectorError, ClientOSError, - ClientPayloadError, ContentTypeError,ServerDisconnectedError) + ClientPayloadError, ContentTypeError, + ServerDisconnectedError) from aiohttp.client_reqrep import ClientResponse from aiohttp_socks import ChainProxyConnector, ProxyConnector, ProxyType from database.models.media_table import media_table @@ -87,6 +88,17 @@ def __init__( self.dynamic_rules = dynamic_rules self.auth = auth + def create_client_session(self): + proxies = self.proxies + proxy = self.proxies[randint(0, len(proxies) - 1)] if proxies else "" + connector = ProxyConnector.from_url(proxy) if proxy else None + + final_cookies = self.auth.cookies if hasattr(self.auth, "cookies") else {} + client_session = ClientSession( + connector=connector, cookies=final_cookies, read_timeout=None + ) + return client_session + def add_sessions(self, original_sessions: list, overwrite_old_sessions=True): if overwrite_old_sessions: sessions = [] @@ -152,12 +164,7 @@ async def json_request( custom_session = False if not session: custom_session = True - proxies = self.proxies - proxy = self.proxies[randint(0, len(proxies) - 1)] if proxies else "" - connector = ProxyConnector.from_url(proxy) if proxy else None - session = ClientSession( - connector=connector, cookies=self.auth.cookies, read_timeout=None - ) + session = self.create_client_session() headers = self.session_rules(link) headers["accept"] = "application/json, text/plain, */*" headers["Connection"] = "keep-alive" @@ -170,41 +177,32 @@ async def json_request( request_method = session.post elif method == "DELETE": request_method = session.delete + result = None while True: try: - async with request_method( - link, headers=headers, data=payload - ) as response: - if method == "HEAD": + response = await request_method(link, headers=headers, data=payload) + if method == "HEAD": + result = response + else: + if json_format and not stream: + result = await response.json() + elif stream and not json_format: result = response else: - if json_format and not stream: - result = await response.json() - elif stream and not json_format: - buffer = [] - if response.status == 200: - async for data in response.content.iter_chunked(4096): - buffer.append(data) - length = len(data) - progress_bar.update(length) - else: - if response.content_length: - progress_bar.update_total_size( - -response.content_length - ) - final_buffer = b"".join(buffer) - buffer.clear() - result = [response, final_buffer] - print - else: - result = await response.read() - if custom_session: - await session.close() - return result + result = await response.read() + break except ClientConnectorError as e: - return - except (ClientPayloadError, ContentTypeError, ClientOSError,ServerDisconnectedError) as e: + break + except ( + ClientPayloadError, + ContentTypeError, + ClientOSError, + ServerDisconnectedError, + ) as e: continue + if custom_session: + await session.close() + return result async def async_requests(self, items: list[str], json_format=True): tasks = [] @@ -256,7 +254,9 @@ async def check(download_item: media_table, response: ClientResponse): filepath = os.path.join( download_item.directory, download_item.filename ) + response_status = False if response.status == 200: + response_status = True if response.content_length: download_item.size = response.content_length @@ -266,7 +266,8 @@ async def check(download_item: media_table, response: ClientResponse): else: return download_item else: - return download_item + if response_status: + return download_item for download_item in download_list: temp_response = [ @@ -288,19 +289,19 @@ async def check(download_item: media_table, response: ClientResponse): [progress_bar.update_total_size(x.size) for x in download_list] async def process_download(download_item: media_table): - response = await self.download_content( + result = await self.download_content( download_item, session, progress_bar, subscription ) - if response: - data, download_item = response.values() - if data: + if result: + response, download_item = result.values() + if response: download_path = os.path.join( download_item.directory, download_item.filename ) - os.makedirs(os.path.dirname(download_path), exist_ok=True) - with open(download_path, "wb") as f: - f.write(data) - download_item.size = len(data) + await main_helper.write_data( + response, download_path, progress_bar + ) + download_item.size = response.content_length download_item.downloaded = True max_threads = calculate_max_threads(self.max_threads) @@ -333,7 +334,7 @@ async def download_content( if not download_item.link: continue response: ClientResponse - response, task = await asyncio.ensure_future( + response = await asyncio.ensure_future( self.json_request( download_item.link, session, @@ -368,7 +369,7 @@ async def download_content( link = main_helper.link_picker(media, quality) download_item.link = link continue - new_task["response"] = task + new_task["response"] = response new_task["download_item"] = download_item break return new_task diff --git a/helpers/main_helper.py b/helpers/main_helper.py index 80e89e59e..f5a699a97 100644 --- a/helpers/main_helper.py +++ b/helpers/main_helper.py @@ -1,6 +1,3 @@ -from database.models.media_table import media_table -from sqlalchemy.orm.session import Session -from apis.onlyfans.classes.extras import content_types import copy import json import math @@ -22,13 +19,17 @@ import psutil import requests import ujson +from aiohttp.client_reqrep import ClientResponse from apis.onlyfans import onlyfans as OnlyFans from apis.onlyfans.classes import create_user +from apis.onlyfans.classes.extras import content_types from bs4 import BeautifulSoup from classes.prepare_metadata import format_variables, prepare_reformat +from database.models.media_table import media_table from mergedeep import Strategy, merge from sqlalchemy import inspect from sqlalchemy.orm import declarative_base +from sqlalchemy.orm.session import Session from tqdm import tqdm import helpers.db_helper as db_helper @@ -726,9 +727,19 @@ def is_me(user_api): return False -async def write_data(download_path:str, data:bytes): - with open(download_path, "wb") as f: - f.write(data) +async def write_data(response: ClientResponse, download_path: str, progress_bar): + if response.status == 200: + os.makedirs(os.path.dirname(download_path), exist_ok=True) + with open(download_path, "wb") as f: + async for data in response.content.iter_chunked(4096): + length = len(data) + progress_bar.update(length) + f.write(data) + return True + else: + if response.content_length: + progress_bar.update_total_size(-response.content_length) + return False def export_data( @@ -906,7 +917,7 @@ def module_chooser(domain, json_sites): return string, site_names -def link_picker(media,video_quality): +def link_picker(media, video_quality): link = "" if "source" in media: quality_key = "source" @@ -929,4 +940,3 @@ def link_picker(media,video_quality): if "src" in media: link = media["src"] return link - \ No newline at end of file diff --git a/modules/onlyfans.py b/modules/onlyfans.py index b2718ca2b..95bc0c897 100644 --- a/modules/onlyfans.py +++ b/modules/onlyfans.py @@ -288,10 +288,10 @@ async def profile_scraper( header = y.header if avatar: override_media_types.append(["Avatars", avatar]) - elif header: override_media_types.append(["Headers", header]) progress_bar = download_session() progress_bar.start(unit="B", unit_scale=True, miniters=1) + session = authed.session_manager.create_client_session() for override_media_type in override_media_types: new_dict = dict() media_type = override_media_type[0] @@ -301,21 +301,20 @@ async def profile_scraper( os.makedirs(directory2, exist_ok=True) download_path = os.path.join(directory2, media_link.split("/")[-2] + ".jpg") response = await authed.session_manager.json_request(media_link, method="HEAD") - if overwrite_files: - if os.path.isfile(download_path): - if os.path.getsize(download_path) == response.content_length: - continue + if os.path.isfile(download_path): + if os.path.getsize(download_path) == response.content_length: + continue progress_bar.update_total_size(response.content_length) - response, data = await authed.session_manager.json_request( + response = await authed.session_manager.json_request( media_link, + session=session, stream=True, json_format=False, sleep=False, progress_bar=progress_bar, ) - downloaded = await main_helper.write_data(download_path, data) - if not downloaded: - continue + downloaded = await main_helper.write_data(response,download_path,progress_bar) + await session.close() progress_bar.close() diff --git a/start_ofd.py b/start_ofd.py index 96fdcd354..a98455ab4 100755 --- a/start_ofd.py +++ b/start_ofd.py @@ -58,4 +58,6 @@ async def main(): input() - asyncio.run(main()) + loop = asyncio.get_event_loop() + loop.run_until_complete(main()) + loop.close()