Этот коммит содержится в:
viginum-datalab 2025-01-15 17:08:16 +01:00
родитель 393cfbc425
Коммит 6431dac7d0
2 изменённых файлов: 26 добавлений и 19 удалений

Просмотреть файл

@ -1,7 +1,9 @@
from contextlib import aclosing from contextlib import aclosing
from httpx import Response from httpx import Response
from json import loads
from typing_extensions import deprecated from typing_extensions import deprecated
from zstd import decompress
from .accounts_pool import AccountsPool from .accounts_pool import AccountsPool
from .logger import set_log_level from .logger import set_log_level
@ -126,7 +128,8 @@ class API:
if rep is None: if rep is None:
return return
obj = rep.json() encoding: str | None = rep.headers.get("content-encoding")
obj = loads(decompress(rep.content)) if encoding == "zstd" else rep.json()
els = get_by_path(obj, "entries") or [] els = get_by_path(obj, "entries") or []
els = [ els = [
x x
@ -138,7 +141,7 @@ class API:
] ]
cur = self._get_cursor(obj, cursor_type) cur = self._get_cursor(obj, cursor_type)
rep, cnt, active = self._is_end(rep, queue, els, cur, cnt, limit) rep, cnt, active = self._is_end(obj, queue, els, cur, cnt, limit)
if rep is None: if rep is None:
return return
@ -169,7 +172,7 @@ class API:
async def search(self, q: str, limit=-1, kv=None): async def search(self, q: str, limit=-1, kv=None):
async with aclosing(self.search_raw(q, limit=limit, kv=kv)) as gen: async with aclosing(self.search_raw(q, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
yield x yield x
# user_by_id # user_by_id
@ -261,7 +264,7 @@ class API:
async def tweet_replies(self, twid: int, limit=-1, kv=None): async def tweet_replies(self, twid: int, limit=-1, kv=None):
async with aclosing(self.tweet_replies_raw(twid, limit=limit, kv=kv)) as gen: async with aclosing(self.tweet_replies_raw(twid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
if x.inReplyToTweetId == twid: if x.inReplyToTweetId == twid:
yield x yield x
@ -278,7 +281,7 @@ class API:
async def followers(self, uid: int, limit=-1, kv=None): async def followers(self, uid: int, limit=-1, kv=None):
async with aclosing(self.followers_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.followers_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# verified_followers # verified_followers
@ -296,7 +299,7 @@ class API:
async def verified_followers(self, uid: int, limit=-1, kv=None): async def verified_followers(self, uid: int, limit=-1, kv=None):
async with aclosing(self.verified_followers_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.verified_followers_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# following # following
@ -311,7 +314,7 @@ class API:
async def following(self, uid: int, limit=-1, kv=None): async def following(self, uid: int, limit=-1, kv=None):
async with aclosing(self.following_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.following_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# subscriptions # subscriptions
@ -326,7 +329,7 @@ class API:
async def subscriptions(self, uid: int, limit=-1, kv=None): async def subscriptions(self, uid: int, limit=-1, kv=None):
async with aclosing(self.subscriptions_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.subscriptions_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# retweeters # retweeters
@ -341,7 +344,7 @@ class API:
async def retweeters(self, twid: int, limit=-1, kv=None): async def retweeters(self, twid: int, limit=-1, kv=None):
async with aclosing(self.retweeters_raw(twid, limit=limit, kv=kv)) as gen: async with aclosing(self.retweeters_raw(twid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# favoriters # favoriters
@ -358,7 +361,7 @@ class API:
async def favoriters(self, twid: int, limit=-1, kv=None): async def favoriters(self, twid: int, limit=-1, kv=None):
async with aclosing(self.favoriters_raw(twid, limit=limit, kv=kv)) as gen: async with aclosing(self.favoriters_raw(twid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_users(rep.json(), limit): for x in parse_users(rep, limit):
yield x yield x
# user_tweets # user_tweets
@ -381,7 +384,7 @@ class API:
async def user_tweets(self, uid: int, limit=-1, kv=None): async def user_tweets(self, uid: int, limit=-1, kv=None):
async with aclosing(self.user_tweets_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.user_tweets_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
yield x yield x
# user_tweets_and_replies # user_tweets_and_replies
@ -404,7 +407,7 @@ class API:
async def user_tweets_and_replies(self, uid: int, limit=-1, kv=None): async def user_tweets_and_replies(self, uid: int, limit=-1, kv=None):
async with aclosing(self.user_tweets_and_replies_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.user_tweets_and_replies_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
yield x yield x
# user_media # user_media
@ -476,7 +479,7 @@ class API:
async def liked_tweets(self, uid: int, limit=-1, kv=None): async def liked_tweets(self, uid: int, limit=-1, kv=None):
async with aclosing(self.liked_tweets_raw(uid, limit=limit, kv=kv)) as gen: async with aclosing(self.liked_tweets_raw(uid, limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
yield x yield x
# Get current user bookmarks # Get current user bookmarks
@ -502,5 +505,5 @@ class API:
async def bookmarks(self, limit=-1, kv=None): async def bookmarks(self, limit=-1, kv=None):
async with aclosing(self.bookmarks_raw(limit=limit, kv=kv)) as gen: async with aclosing(self.bookmarks_raw(limit=limit, kv=kv)) as gen:
async for rep in gen: async for rep in gen:
for x in parse_tweets(rep.json(), limit): for x in parse_tweets(rep, limit):
yield x yield x

Просмотреть файл

@ -1,6 +1,7 @@
import json from json import JSONDecodeError, dumps, loads
import os import os
from typing import Any from typing import Any
from zstd import decompress
import httpx import httpx
from httpx import AsyncClient, Response from httpx import AsyncClient, Response
@ -56,8 +57,10 @@ def dump_rep(rep: Response):
msg.append("\n") msg.append("\n")
try: try:
msg.append(json.dumps(rep.json(), indent=2)) encoding: str | None = rep.headers.get("content-encoding")
except json.JSONDecodeError: obj = loads(decompress(rep.content)) if encoding == "zstd" else rep.json()
msg.append(dumps(obj, indent=2))
except JSONDecodeError:
msg.append(rep.text) msg.append(rep.text)
txt = "\n".join(msg) txt = "\n".join(msg)
@ -120,8 +123,9 @@ class QueueClient:
dump_rep(rep) dump_rep(rep)
try: try:
res = rep.json() encoding: str | None = rep.headers.get("content-encoding")
except json.JSONDecodeError: res = loads(decompress(rep.content)) if encoding == "zstd" else rep.json()
except JSONDecodeError:
res: Any = {"_raw": rep.text} res: Any = {"_raw": rep.text}
limit_remaining = int(rep.headers.get("x-rate-limit-remaining", -1)) limit_remaining = int(rep.headers.get("x-rate-limit-remaining", -1))