Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle zstd encoding #230

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Handle zstd encoding #230

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
393cfbc
build: add zstd to dependencies
viginum-datalab Jan 15, 2025
6431dac
fix: decode zstd response
viginum-datalab Jan 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ dev = [
"pytest-httpx>=0.28.0",
"pytest>=7.4.4",
"ruff>=0.1.11",
"zstd>=1.5.6.1",
]

[project.urls]
Expand Down
31 changes: 17 additions & 14 deletions twscrape/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
from contextlib import aclosing

from httpx import Response
from json import loads
from typing_extensions import deprecated
from zstd import decompress

from .accounts_pool import AccountsPool
from .logger import set_log_level
Expand Down Expand Up @@ -126,7 +128,8 @@ async def _gql_items(
if rep is None:
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 = [
x
Expand All @@ -138,7 +141,7 @@ async def _gql_items(
]
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:
return

Expand Down Expand Up @@ -169,7 +172,7 @@ async def search_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
yield x

# user_by_id
Expand Down Expand Up @@ -261,7 +264,7 @@ async def tweet_replies_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
if x.inReplyToTweetId == twid:
yield x

Expand All @@ -278,7 +281,7 @@ async def followers_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# verified_followers
Expand All @@ -296,7 +299,7 @@ async def verified_followers_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# following
Expand All @@ -311,7 +314,7 @@ async def following_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# subscriptions
Expand All @@ -326,7 +329,7 @@ async def subscriptions_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# retweeters
Expand All @@ -341,7 +344,7 @@ async def retweeters_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# favoriters
Expand All @@ -358,7 +361,7 @@ async def favoriters_raw(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 for rep in gen:
for x in parse_users(rep.json(), limit):
for x in parse_users(rep, limit):
yield x

# user_tweets
Expand All @@ -381,7 +384,7 @@ async def user_tweets_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
yield x

# user_tweets_and_replies
Expand All @@ -404,7 +407,7 @@ async def user_tweets_and_replies_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
yield x

# user_media
Expand Down Expand Up @@ -476,7 +479,7 @@ async def liked_tweets_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
yield x

# Get current user bookmarks
Expand All @@ -502,5 +505,5 @@ async def bookmarks_raw(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 for rep in gen:
for x in parse_tweets(rep.json(), limit):
for x in parse_tweets(rep, limit):
yield x
14 changes: 9 additions & 5 deletions twscrape/queue_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
from json import JSONDecodeError, dumps, loads
import os
from typing import Any
from zstd import decompress

import httpx
from httpx import AsyncClient, Response
Expand Down Expand Up @@ -56,8 +57,10 @@ def dump_rep(rep: Response):
msg.append("\n")

try:
msg.append(json.dumps(rep.json(), indent=2))
except json.JSONDecodeError:
encoding: str | None = rep.headers.get("content-encoding")
obj = loads(decompress(rep.content)) if encoding == "zstd" else rep.json()
msg.append(dumps(obj, indent=2))
except JSONDecodeError:
msg.append(rep.text)

txt = "\n".join(msg)
Expand Down Expand Up @@ -120,8 +123,9 @@ async def _check_rep(self, rep: Response) -> None:
dump_rep(rep)

try:
res = rep.json()
except json.JSONDecodeError:
encoding: str | None = rep.headers.get("content-encoding")
res = loads(decompress(rep.content)) if encoding == "zstd" else rep.json()
except JSONDecodeError:
res: Any = {"_raw": rep.text}

limit_remaining = int(rep.headers.get("x-rate-limit-remaining", -1))
Expand Down