Merge pull request #48 from atcoder/patch/refactor

refactor expander

Author: yosupo06

expander.py

@@ -6,40 +6,69 @@
 from logging import Logger, basicConfig, getLogger
 from os import getenv, environ
 from pathlib import Path
-from typing import List
+from typing import List, Set, Optional
 
 
 logger = getLogger(__name__)  # type: Logger
 
-atcoder_include = re.compile('#include\s*["<](atcoder/[a-z_]*(|.hpp))[">]\s*')
 
-include_guard = re.compile('#.*ATCODER_[A-Z_]*_HPP')
-
-lib_path = Path.cwd()
-
-defined = set()
-
-def dfs(f: str) -> List[str]:
-    global defined
-    if f in defined:
-        logger.info('already included {}, skip'.format(f))
-        return []
-    defined.add(f)
-
-    logger.info('include {}'.format(f))
-
-    s = open(str(lib_path / f)).read()
-    result = []
-    for line in s.splitlines():
-        if include_guard.match(line):
-            continue
-
-        m = atcoder_include.match(line)
-        if m:
-            result.extend(dfs(m.group(1)))
-            continue
-        result.append(line)
-    return result
+class Expander:
+    atcoder_include = re.compile(
+        '#include\s*["<](atcoder/[a-z_]*(|.hpp))[">]\s*')
+
+    include_guard = re.compile('#.*ATCODER_[A-Z_]*_HPP')
+
+    def __init__(self, lib_paths: List[Path] = None):
+        if lib_paths:
+            self.lib_paths = lib_paths
+        else:
+            self.lib_paths = [Path.cwd()]
+
+    included = set()  # type: Set[str]
+
+    def find_acl(self, acl_name: str) -> Optional[Path]:
+        for lib_path in self.lib_paths:
+            path = lib_path / acl_name
+            if path.exists():
+                return path
+        return None
+
+    def expand_acl(self, acl_name: str) -> List[str]:
+        if acl_name in self.included:
+            logger.info('already included: {}'.format(acl_name))
+            return []
+        self.included.add(acl_name)
+        logger.info('include: {}'.format(acl_name))
+        acl_path = self.find_acl(acl_name)
+        if not acl_path:
+            logger.warning('cannot find: {}'.format(acl_name))
+            raise FileNotFoundError()
+
+        acl_source = open(str(acl_path)).read()
+
+        result = []  # type: List[str]
+        for line in acl_source.splitlines():
+            if self.include_guard.match(line):
+                continue
+
+            m = self.atcoder_include.match(line)
+            if m:
+                result.extend(self.expand_acl(m.group(1)))
+                continue
+            result.append(line)
+        return result
+
+    def expand(self, source: str) -> str:
+        self.included = set()
+        result = []  # type: List[str]
+        for line in source.splitlines():
+            m = self.atcoder_include.match(line)
+
+            if m:
+                result.extend(self.expand_acl(m.group(1)))
+                continue
+            result.append(line)
+        return '\n'.join(result)
 
 
 if __name__ == "__main__":
@@ -55,22 +84,16 @@ def dfs(f: str) -> List[str]:
     parser.add_argument('--lib', help='Path to Atcoder Library')
     opts = parser.parse_args()
 
+    lib_path = Path.cwd()
     if opts.lib:
         lib_path = Path(opts.lib)
     elif 'CPLUS_INCLUDE_PATH' in environ:
         lib_path = Path(environ['CPLUS_INCLUDE_PATH'])
-    s = open(opts.source).read()
-
-    result = []
-    for line in s.splitlines():
-        m = atcoder_include.match(line)
 
-        if m:
-            result.extend(dfs(m.group(1)))
-            continue
-        result.append(line)
+    expander = Expander([lib_path])
+    source = open(opts.source).read()
+    output = expander.expand(source)
 
-    output = '\n'.join(result) + '\n'
     if opts.console:
         print(output)
     else:

test/expander/include_unionfind.cpp renamed to test/expander/include_dsu.cpp

@@ -13,30 +13,34 @@
 
 
 class Test(unittest.TestCase):
-    def compile_test(self, source: Path, env=None):
+    def compile_test(self, source: Path, expander_args=[], env=None):
         if not env:
             env = environ.copy()
         lib_dir = Path.cwd().resolve()
         expander_path = Path('expander.py').resolve()
         with TemporaryDirectory() as new_dir:
             tmp = Path(new_dir)
             proc = run(['python', str(expander_path), str(
-                source.resolve()), '--lib', str(lib_dir)], cwd=str(tmp), env=env)
+                source.resolve())] + expander_args, cwd=str(tmp), env=env)
             self.assertEqual(proc.returncode, 0)
             proc = run(['g++', 'combined.cpp', '-std=c++14'], cwd=str(tmp))
             self.assertEqual(proc.returncode, 0)
 
-    def test_unionfind(self):
-        self.compile_test(Path('test/expander/include_unionfind.cpp'))
+    def test_dsu(self):
+        self.compile_test(Path('test/expander/include_dsu.cpp'),
+                          expander_args=['--lib', str(Path.cwd().resolve())])
 
     def test_unusual_format(self):
-        self.compile_test(Path('test/expander/include_unusual_format.cpp'))
+        self.compile_test(Path('test/expander/include_unusual_format.cpp'),
+                          expander_args=['--lib', str(Path.cwd().resolve())])
 
     def test_all(self):
-        self.compile_test(Path('test/expander/include_all.cpp'))
+        self.compile_test(Path('test/expander/include_all.cpp'),
+                          expander_args=['--lib', str(Path.cwd().resolve())])
 
     def test_comment_out(self):
-        self.compile_test(Path('test/expander/comment_out.cpp'))
+        self.compile_test(Path('test/expander/comment_out.cpp'),
+                          expander_args=['--lib', str(Path.cwd().resolve())])
 
     def test_env_value(self):
         env = environ.copy()