feat(copy): support binary source

JinnLynn · Jul 22, 2024 · 2ee003b · 2ee003b
1 parent 5034aae
commit 2ee003b
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 12 deletions.
diff --git a/src/genpac/core.py b/src/genpac/core.py
@@ -324,11 +324,21 @@ def generate(self):
 
         content = self.formater.generate(replacements)
 
+        if content is None:
+            raise FatalError(f'{self.formater._name}: 生成失败')
+
         output = self.options.output
         if not output or output == '-':
-            sys.stdout.write(content)
+            if isinstance(content, str):
+                sys.stdout.write(content)
+            else:
+                logger.error('无法写入bytes内容到stdout.')
         else:
-            write_file(output, content, fail_msg='写入输出文件`{path}`失败')
+            try:
+                with open(output, 'w' if isinstance(content, str) else 'wb') as fp:
+                    fp.write(content)
+            except Exception as e:
+                raise FatalError(f'写入输出文件`{output}`失败: {e}')
 
         self.formater.post_generate()
 
@@ -392,17 +402,21 @@ def request(self, url):
     def fetch_online(self, url):
         try:
             content = self.request(url)
+            if content:
+                self.__class__._cache[url] = content
         except Exception as e:
             logger.error(f'Fetch online fail: {e} {url}')
             return
-        return content.decode() if isinstance(content, bytes) else content
+        return content
 
     # 使用类变量缓存在线获取的内容
-    def fetch(self, url):
+    def fetch(self, url, decode=True):
         content = self.__class__._cache.get(url) or self.fetch_online(url)
+
         if content:
-            self.__class__._cache[url] = content
-        return content
+            if decode in (False, None):
+                return content
+            return content.decode('utf-8' if not isinstance(decode, str) else decode)
 
     def fetch_gfwlist(self):
         if self.options.gfwlist_disabled or self.options.gfwlist_url == '-':

diff --git a/src/genpac/format/base.py b/src/genpac/format/base.py
@@ -84,8 +84,8 @@ def error(self, msg):
     def replace(self, text, replacements, **kwargs):
         return replace_all(text, replacements, **kwargs)
 
-    def fetch(self, url):
-        return self.generator.fetch(url)
+    def fetch(self, url, decode=True):
+        return self.generator.fetch(url, decode=decode)
 
     # RETURN:
     # [0][0]: user.direct    [0][1]: user.proxy

diff --git a/src/genpac/format/copy.py b/src/genpac/format/copy.py
@@ -1,7 +1,7 @@
 from os import path
 
 from .base import formater, FmtBase
-from ..util import open_file, FatalError
+from ..util import FatalError
 
 
 @formater('copy', desc="IP地址列表")
@@ -17,13 +17,13 @@ def prepare(cls, parser):
         cls.register_option('source', metavar='SRC', help='来源, 网址或文件路径')
 
     def generate(self, replacements):
-        content = ''
+        content = None
         try:
             if path.isfile(self.options.source):
-                with open_file(self.options.source) as fp:
+                with open(self.options.source, 'rb') as fp:
                     content = fp.read()
             else:
-                content = self.fetch(self.options.source)
+                content = self.fetch(self.options.source, decode=False)
                 if content is None:
                     raise Exception
         except Exception: