mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-09-21 12:42:22 +08:00
icecraftandshenguanlin
40e0827e60
* feat: refractor cli command * feat: add docs to describe the output files of cli * feat: resove review comments * feat: updat docs about middle.json --------- Co-authored-by: shenguanlin <shenguanlin@pjlab.org.cn>
75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
import os
|
|
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
|
|
from loguru import logger
|
|
|
|
|
|
class DiskReaderWriter(AbsReaderWriter):
|
|
def __init__(self, parent_path, encoding="utf-8"):
|
|
self.path = parent_path
|
|
self.encoding = encoding
|
|
|
|
def read(self, path, mode=AbsReaderWriter.MODE_TXT):
|
|
if os.path.isabs(path):
|
|
abspath = path
|
|
else:
|
|
abspath = os.path.join(self.path, path)
|
|
if not os.path.exists(abspath):
|
|
logger.error(f"file {abspath} not exists")
|
|
raise Exception(f"file {abspath} no exists")
|
|
if mode == AbsReaderWriter.MODE_TXT:
|
|
with open(abspath, "r", encoding=self.encoding) as f:
|
|
return f.read()
|
|
elif mode == AbsReaderWriter.MODE_BIN:
|
|
with open(abspath, "rb") as f:
|
|
return f.read()
|
|
else:
|
|
raise ValueError("Invalid mode. Use 'text' or 'binary'.")
|
|
|
|
def write(self, content, path, mode=AbsReaderWriter.MODE_TXT):
|
|
if os.path.isabs(path):
|
|
abspath = path
|
|
else:
|
|
abspath = os.path.join(self.path, path)
|
|
directory_path = os.path.dirname(abspath)
|
|
if not os.path.exists(directory_path):
|
|
os.makedirs(directory_path)
|
|
if mode == AbsReaderWriter.MODE_TXT:
|
|
with open(abspath, "w", encoding=self.encoding, errors="replace") as f:
|
|
f.write(content)
|
|
|
|
elif mode == AbsReaderWriter.MODE_BIN:
|
|
with open(abspath, "wb") as f:
|
|
f.write(content)
|
|
else:
|
|
raise ValueError("Invalid mode. Use 'text' or 'binary'.")
|
|
|
|
def read_offset(self, path: str, offset=0, limit=None):
|
|
abspath = path
|
|
if not os.path.isabs(path):
|
|
abspath = os.path.join(self.path, path)
|
|
with open(abspath, "rb") as f:
|
|
f.seek(offset)
|
|
return f.read(limit)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if 0:
|
|
file_path = "io/test/example.txt"
|
|
drw = DiskReaderWriter("D:\projects\papayfork\Magic-PDF\magic_pdf")
|
|
|
|
# 写入内容到文件
|
|
drw.write(b"Hello, World!", path="io/test/example.txt", mode="binary")
|
|
|
|
# 从文件读取内容
|
|
content = drw.read(path=file_path)
|
|
if content:
|
|
logger.info(f"从 {file_path} 读取的内容: {content}")
|
|
if 1:
|
|
drw = DiskReaderWriter("/opt/data/pdf/resources/test/io/")
|
|
content_bin = drw.read_offset("1.txt")
|
|
assert content_bin == b"ABCD!"
|
|
|
|
content_bin = drw.read_offset("1.txt", offset=1, limit=2)
|
|
assert content_bin == b"BC"
|
|
|