async def atomic_write(
path: Union[Path, str],
mode: str = "wt",
**kwargs,
) -> AsyncIterator[Union[AsyncTextIOWrapper, AsyncBufferedReader]]:
final = AsyncPath(path)
prefix = f".{final.stem}."
suffix = f"{final.suffix}.partial"
async with aiofiles.tempfile.NamedTemporaryFile(
dir=final.parent,
prefix=prefix,
suffix=suffix,
delete=False,
) as file:
temp = AsyncPath(file.name)
try:
if mode not in TRUNCATE_FILE_MODES and await final.exists():
await copy_file_with_metadata(final, temp)
async with aiofiles.open(temp, mode, **kwargs) as out: # type: ignore
yield out
await temp.replace(final)
finally:
with suppress(FileNotFoundError):
await temp.unlink()
async def file_transition_as(file: Path, root_folder: Path, dist: str):
a_file = AsyncPath(file)
a_root_folder = AsyncPath(root_folder)
target_folder = a_root_folder / dist
await target_folder.mkdir(exist_ok=True)
ext = AsyncPath(file).suffix
new_name = normalize(a_file.name.replace(ext, '')) + ext
await a_file.replace(target_folder / new_name)
async def test_readme_example5_glob():
home: AsyncPath = await AsyncPath.home()
async for path in home.glob(WILDCARD_GLOB):
assert isinstance(path, AsyncPath)
pkg_dir = AsyncPath(__file__).parent
assert await pkg_dir.exists()
paths = [path async for path in pkg_dir.glob(RECURSIVE_GLOB)]
assert len(paths) > NO_PATHS
async def uploadFolder(
self,
sourceFolder: AsyncPath,
*,
gid: Optional[str] = None,
parent_id: Optional[str] = None,
msg: Optional[pyrogram.types.Message] = None
) -> AsyncIterator[asyncio.Task[None]]:
async for content in sourceFolder.iterdir():
if await content.is_dir():
childFolder = await self.createFolder(content.name, parent_id)
async for task in self.uploadFolder(content,
gid=gid,
parent_id=childFolder,
msg=msg):
yield task
elif await content.is_file():
file = util.File(content)
content = await self.uploadFile(file, parent_id, msg)
if isinstance(content, str): # Skip because file size is 0
continue
yield self.bot.loop.create_task(file.progress(update=False),
name=gid)
await asyncio.sleep(0.5)
else:
raise ValueError(f"{content} is not a file or folder")
def _test_is_pure( path: Path, apath: AsyncPath, ): # AsyncPurePath methods are not async assert str(path) == str(apath) assert path.anchor == apath.anchor assert path.name == apath.name assert path.drive == apath.drive assert path.parts == apath.parts assert path.root == apath.root assert path.stem == apath.stem assert path.suffix == apath.suffix assert path.suffixes == apath.suffixes assert path.as_uri() == apath.as_uri() assert path.as_posix() == apath.as_posix() assert path.is_absolute() == apath.is_absolute() assert path.is_reserved() == apath.is_reserved() assert path.is_relative_to(path) == apath.is_relative_to(apath)
def __init__(self) -> None:
config: MutableMapping[_KT, Any] = {
"api_id": os.environ.get("API_ID"),
"api_hash": os.environ.get("API_HASH"),
"bot_token": os.environ.get("BOT_TOKEN"),
"container": os.environ.get("CONTAINER") == "True",
"db_uri": os.environ.get("DB_URI"),
"db_uri_anjani": os.environ.get("DB_URI_ANJANI"),
"download_path": os.environ.get("DOWNLOAD_PATH"),
"gdrive_folder_id": os.environ.get("G_DRIVE_FOLDER_ID"),
"gdrive_index_link": os.environ.get("G_DRIVE_INDEX_LINK"),
"gdrive_secret": os.environ.get("G_DRIVE_SECRET"),
"github_repo": os.environ.get("GITHUB_REPO"),
"github_token": os.environ.get("GITHUB_TOKEN"),
"heroku_api_key": os.environ.get("HEROKU_API_KEY"),
"heroku_app_name": os.environ.get("HEROKU_APP"),
"mirror_enabled": os.environ.get("MIRROR_MODULE") == "enable",
"sp_token": os.environ.get("SP_TOKEN"),
"sp_url": os.environ.get("SP_URL"),
"string_session": os.environ.get("STRING_SESSION"),
}
for key, value in config.items():
if not value:
if key == "download_path":
value = AsyncPath(Path.home() / "downloads")
elif key == "github_repo":
value = "adekmaulana/caligo"
if value == "":
value = None
else:
if key == "download_path":
value = AsyncPath(value)
elif key == "gdrive_index_link":
value = value.rstrip("/")
elif key == "gdrive_secret":
value = json.loads(value)
super().__setattr__(key, value)
self.__data[key] = value
def name(self) -> str:
if not self._name:
filePath = str(self._path.absolute())
dirPath = str(self._path.parent.absolute())
if filePath.startswith(dirPath):
start = len(dirPath) + 1
self._name = AsyncPath(filePath[start:]).parts[0]
else:
self._name = self._path.parts[-1]
return self._name
async def order(self, keywords, queue):
kw = " ".join(keywords)
if len(kw) < 3:
return
cmds = []
async for path in AsyncPath().rglob("*"):
if keywords[0] in str(path):
cmds.append(str(path))
if len(cmds) > 100:
self._put_cmd(kw, cmds, queue)
cmds = []
self._put_cmd(kw, cmds, queue)
async def test_readme_example4_read_write():
text: str = 'example'
async with NamedTemporaryFile() as temp:
path = AsyncPath(temp.name)
async with path.open(mode='w') as file:
await file.write(text)
async with path.open(mode='r') as file:
result: str = await file.read()
assert result == text
async with NamedTemporaryFile() as temp:
path = AsyncPath(temp.name)
await path.write_text(text)
result: str = await path.read_text()
assert result == text
content: bytes = text.encode()
await path.write_bytes(content)
result: bytes = await path.read_bytes()
assert result == content
async def test_readme_example2_convert():
home: Path = Path.home()
ahome: AsyncPath = AsyncPath(home)
path: Path = Path(ahome)
assert isinstance(home, Path)
assert isinstance(ahome, AsyncPath)
assert isinstance(path, Path)
# AsyncPath and Path objects can point to the same file
assert str(home) == str(ahome) == str(path)
# but AsyncPath and Path objects are not equivalent
assert not home == ahome
async def create_tasks(
data_path: AsyncPath, root_path: Path, aiosession: ClientSession,
logger: loguru.Logger
) -> tuple[dict[str, Any], AsyncIterable[tuple[AuditTask, LogQueue]]]:
async with data_path.open('r') as data_f:
meta_line = await anext(aiter(data_f))
audit_meta = json_parse_str(meta_line)
if audit_meta['state'] != 'meta':
raise ValueError(
'Expected first line of task list to be meta record')
return audit_meta, _create_task_stream(
data_path,
pendulum.from_timestamp(audit_meta['max_last_modified']),
root_path, aiosession, logger)
async def _create_task_stream(
task_path: AsyncPath, max_last_modified: DateTime, root_path: Path,
aiosession: ClientSession,
logger: loguru.Logger) -> AsyncIterable[tuple[AuditTask, LogQueue]]:
async with aiofiles.open(task_path, 'r') as task_f:
_ = await task_f.readline() # Skip over meta line
while (lines := await task_f.readlines(2**26)):
for data_line in lines:
mb_data = json_parse_str(data_line)
task_path = AsyncPath(_fanout_path(root_path, mb_data['id']))
log_queue = LogQueue(task_path / 'audit_log')
task_logger = logger.bind(log_queue=log_queue)
yield AuditTask(mb_data, max_last_modified, task_path,
aiosession, task_logger), log_queue
async def do_audit(mb_data_file_path: Path, output_path: Path,
concurrency: int, spam: bool) -> None:
configure_logging(spam)
ia_creds = get_ia_credentials()
if ia_creds is None:
loguru.logger.error('IA credentials not found in ~/.ia')
return
s3_access, s3_secret = ia_creds
with mb_data_file_path.open('r') as f:
num_items = sum(1 for l in f)
task_q: asyncio.Queue[tuple[AuditTask,
LogQueue]] = asyncio.Queue(concurrency * 2)
session = ClientSession(
connector=TCPConnector(limit=concurrency),
headers={'Authorization': f'LOW {s3_access}:{s3_secret}'})
async with session:
audit_meta, task_stream = await create_tasks(
AsyncPath(mb_data_file_path), output_path, session, loguru.logger)
num_items = audit_meta['count']
with ProgressBar(num_items) as progress:
queuer = asyncio.create_task(
queue_tasks(task_stream, task_q, progress))
aggregator = ResultAggregator(output_path, progress)
runners = [
asyncio.create_task(task_runner(task_q, aggregator, progress))
for _ in range(concurrency)
]
# Wait until all tasks are queued
await queuer
# Wait until all tasks are done
await task_q.join()
# Terminate the runners, all jobs done
for runner in runners:
runner.cancel()
aggregator.finish()
write_logs(output_path, aggregator)
write_failed_items(output_path, aggregator)
write_tables(output_path, aggregator)
async def test_bare_init(alice: Client, tmp_path: Path):
client = Client(
tmp_path,
alice.server,
alice.device_id,
alice.user_id,
alice.access_token,
)
assert client.path == AsyncPath(tmp_path / "client.json")
assert read_json(client.path) == {
"server": alice.server,
"user_id": alice.user_id,
"access_token": alice.access_token,
"device_id": alice.device_id,
}
async def from_file_to_move(
cls,
store: "MediaStore",
path: Union[Path, str],
) -> "StoreMedia":
apath = AsyncPath(path)
async with aiofiles.open(apath, "rb") as file:
sha256 = await sha256_chunked(file)
content = store._content_path(sha256)
await content.parent.mkdir(parents=True, exist_ok=True)
await apath.replace(content)
await remove_write_permissions(content)
return cls(store, sha256)
def name(self) -> str:
if not self._name:
if self.bittorrent and self.bittorrent.info:
self._name = self.bittorrent.info["name"]
elif self.files[0].metadata:
self._name = str(self.files[0].path)
else:
file_path = str(self.files[0].path.absolute())
dir_path = str(self.dir.absolute())
if file_path.startswith(dir_path):
start_pos = len(dir_path) + 1
self._name = AsyncPath(file_path[start_pos:]).parts[0]
else:
try:
self._name = self.files[0].uris[0]["uri"].split(
"/")[-1]
except IndexError:
pass
return self._name
def __init__(
self,
task_record: dict[str, Any],
max_last_modified: DateTime,
audit_path: Path,
session: ClientSession,
logger: Logger,
) -> None:
self._record = task_record
self._mbid = task_record['id']
self._max_last_modified = max_last_modified
self.audit_path = AsyncPath(audit_path)
self._logger = logger
self._ia_item = IAItem(f'mbid-{self._mbid}', audit_path, session,
logger)
self._results: list[CheckResult] = []
self._start_stage_cbs = []
self._finish_stage_cbs = []
async def init(cls, bot: Any, drive: Any) -> "Aria2WebSocketServer":
self = cls(bot, drive)
download_path = self.bot.getConfig["download_path"]
await download_path.mkdir(parents=True, exist_ok=True)
link = "https://raw.githubusercontent.com/ngosang/trackerslist/master/trackers_best.txt"
async with self.bot.http.get(link) as resp:
trackers_list: str = await resp.text()
trackers: str = "[" + trackers_list.replace("\n\n", ",") + "]"
cmd = [
"aria2c", f"--dir={str(download_path)}", "--enable-rpc",
"--rpc-listen-all=false", "--max-connection-per-server=10",
"--rpc-max-request-size=1024M", "--seed-time=0.01",
"--seed-ratio=0.1", "--max-concurrent-downloads=5",
"--min-split-size=10M", "--follow-torrent=mem", "--split=10",
"--bt-save-metadata=true", f"--bt-tracker={trackers}",
"--daemon=true", "--allow-overwrite=true"
]
key_path = AsyncPath(Path.home() / ".cache" / "caligo" / ".certs")
if await (key_path / "cert.pem").is_file() and (await
(key_path /
"key.pem").is_file()):
cmd.insert(4, "--rpc-listen-port=8443")
cmd.insert(3, "--rpc-secure=true")
cmd.insert(3, "--rpc-private-key=" + str(key_path / "key.pem"))
cmd.insert(3, "--rpc-certificate=" + str(key_path / "cert.pem"))
self._protocol = "https://localhost:8443/jsonrpc"
else:
cmd.insert(4, "--rpc-listen-port=8100")
self._protocol = "http://127.0.0.1:8100/jsonrpc"
server = AsyncAria2Server(*cmd, daemon=True)
await server.start()
await server.wait()
return self
async def downloadFile(
self, ctx: command.Context, msg: pyrogram.types.Message
) -> Optional[AsyncPath]:
download_path = self.bot.getConfig["download_path"]
before = util.time.sec()
last_update_time = None
human = util.misc.human_readable_bytes
time = util.time.format_duration_td
if msg.document:
file_name = msg.document.file_name
elif msg.audio:
file_name = msg.audio.file_name
elif msg.video:
file_name = msg.video.file_name
elif msg.sticker:
file_name = msg.sticker.file_name
elif msg.photo:
date = datetime.fromtimestamp(msg.photo.date)
file_name = f"photo_{date.strftime('%Y-%m-%d_%H-%M-%S')}.jpg"
elif msg.voice:
date = datetime.fromtimestamp(msg.voice.date)
file_name = f"audio_{date.strftime('%Y-%m-%d_%H-%M-%S')}.ogg"
else:
file_name = None
def prog_func(current: int, total: int) -> None:
nonlocal last_update_time
percent = current / total
after = util.time.sec() - before
now = datetime.now()
try:
speed = round(current / after, 2)
eta = timedelta(seconds=int(round((total - current) / speed)))
except ZeroDivisionError:
speed = 0
eta = timedelta(seconds=0)
bullets = "●" * int(round(percent * 10)) + "○"
if len(bullets) > 10:
bullets = bullets.replace("○", "")
space = ' ' * (10 - len(bullets))
progress = (
f"`{file_name}`\n"
f"Status: **Downloading**\n"
f"Progress: [{bullets + space}] {round(percent * 100)}%\n"
f"__{human(current)} of {human(total)} @ "
f"{human(speed, postfix='/s')}\neta - {time(eta)}__\n\n")
# Only edit message once every 5 seconds to avoid ratelimits
if last_update_time is None or (
now - last_update_time).total_seconds() >= 5:
self.bot.loop.create_task(ctx.respond(progress))
last_update_time = now
file_path = download_path / file_name
file_path = await ctx.bot.client.download_media(msg,
file_name=file_path,
progress=prog_func)
return AsyncPath(file_path) if file_path is not None else file_path
def get_paths(path: PathTypes) -> Paths: return Path(path), AsyncPath(path)
async def extract_backups(path: Path):
await TARGET_DIR.mkdir(exist_ok=True, parents=True)
await asyncio.wait([extract(p) async for p in path.iterdir()])
async def add_write_permissions(path: Union[Path, str]) -> AsyncPath:
apath = AsyncPath(path)
mode = (await apath.stat()).st_mode
read_bits = stat.S_IREAD | stat.S_IRGRP | stat.S_IROTH
await apath.chmod(mode | (mode & read_bits) >> 1) # Copy read → write bits
return apath
async def remove_write_permissions(path: Union[Path, str]) -> AsyncPath:
apath = AsyncPath(path)
mode = (await apath.stat()).st_mode
ro_mask = 0o777 ^ (stat.S_IWRITE | stat.S_IWGRP | stat.S_IWOTH)
await apath.chmod(mode & ro_mask)
return apath
def path(self) -> AsyncPath:
return AsyncPath(self.base_dir) / "client.json"
async def delete_folder_as(folder: Path):
a_folder = AsyncPath(folder)
await a_folder.rmdir()
async def download_file(
ctx: command.Context,
msg: pyrogram.types.Message,
text: Optional[bool] = False
) -> Optional[Union[AsyncPath, str, bytes]]:
"""Downloads the file embedded in the given message."""
download_path = ctx.bot.getConfig["download_path"]
if text is True:
path = await ctx.bot.client.download_media(msg)
if path:
path = AsyncPath(path)
async with aiofile.async_open(path, "r") as file:
content = await file.read()
await path.unlink()
return content
before = sec()
last_update_time = None
if msg.document:
file_name = msg.document.file_name
elif msg.audio:
file_name = msg.audio.file_name
elif msg.video:
file_name = msg.video.file_name
elif msg.sticker:
file_name = msg.sticker.file_name
elif msg.photo:
date = datetime.fromtimestamp(msg.photo.date)
file_name = f"photo_{date.strftime('%Y-%m-%d_%H-%M-%S')}.jpg"
elif msg.voice:
date = datetime.fromtimestamp(msg.voice.date)
file_name = f"audio_{date.strftime('%Y-%m-%d_%H-%M-%S')}.ogg"
else:
file_name = "Unknown"
loop = asyncio.get_event_loop()
def prog_func(current: int, total: int) -> None:
nonlocal last_update_time
if not ctx:
return
percent = current / total
after = sec() - before
now = datetime.now()
try:
speed = round(current / after, 2)
eta = timedelta(seconds=int(round((total - current) / speed)))
except ZeroDivisionError:
speed = 0
eta = timedelta(seconds=0)
bullets = "●" * int(round(percent * 10)) + "○"
if len(bullets) > 10:
bullets = bullets.replace("○", "")
space = ' ' * (10 - len(bullets))
progress = (f"`{file_name}`\n"
f"Status: **Downloading**\n"
f"Progress: [{bullets + space}] {round(percent * 100)}%\n"
f"__{human(current)} of {human(total)} @ "
f"{human(speed, postfix='/s')}\neta - {time(eta)}__\n\n")
# Only edit message once every 5 seconds to avoid ratelimits
if last_update_time is None or (now -
last_update_time).total_seconds() >= 5:
loop.create_task(ctx.respond(progress))
last_update_time = now
path = await ctx.bot.client.download_media(msg,
file_name=str(download_path) +
"/" + file_name,
progress=prog_func)
return AsyncPath(path) if path is not None else path
def dir(self) -> AsyncPath:
return AsyncPath(self._data["dir"])
async def cmd_upload(self, ctx: command.Context) -> Optional[str]:
if not ctx.input:
return "__Pass the file path.__"
before = util.time.sec()
file_path = AsyncPath(ctx.input)
last_update_time = None
if await file_path.is_dir():
await ctx.respond("__The path you input is a directory.__")
return
if not await file_path.is_file():
await ctx.respond("__The file you input doesn't exists.__")
return
human = util.misc.human_readable_bytes
time = util.time.format_duration_td
def prog_func(current: int, total: int) -> None:
nonlocal last_update_time
percent = current / total
after = util.time.sec() - before
now = datetime.now()
try:
speed = round(current / after, 2)
eta = timedelta(seconds=int(round((total - current) / speed)))
except ZeroDivisionError:
speed = 0
eta = timedelta(seconds=0)
bullets = "●" * int(round(percent * 10)) + "○"
if len(bullets) > 10:
bullets = bullets.replace("○", "")
space = ' ' * (10 - len(bullets))
progress = (
f"`{file_path.name}`\n"
f"Status: **Uploading**\n"
f"Progress: [{bullets + space}] {round(percent * 100)}%\n"
f"__{human(current)} of {human(total)} @ "
f"{human(speed, postfix='/s')}\neta - {time(eta)}__\n\n")
# Only edit message once every 5 seconds to avoid ratelimits
if last_update_time is None or (
now - last_update_time).total_seconds() >= 5:
self.bot.loop.create_task(ctx.respond(progress))
last_update_time = now
task = self.bot.loop.create_task(
self.bot.client.send_document(ctx.msg.chat.id,
str(file_path),
force_document=True,
progress=prog_func))
self.task.add((ctx.msg.message_id, task))
try:
await task
except asyncio.CancelledError:
return "__Transmission aborted.__"
else:
self.task.remove((ctx.msg.message_id, task))
await ctx.msg.delete()
return
def __init__(self, identifier: str, cache_dir_path: Path,
session: aiohttp.ClientSession, logger: Logger) -> None:
self._cache_dir_path = AsyncPath(cache_dir_path)
self._identifier = identifier
self._session = session
self._logger = logger
def path(self) -> AsyncPath:
return AsyncPath(self._data["path"])