Read .tofu files

Author: dflook

image/src/terraform/module.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import os
-from typing import Any, cast, NewType, Optional, TYPE_CHECKING, TypedDict, List
+from typing import Any, cast, NewType, Optional, TYPE_CHECKING, TypedDict, List, Iterable
 
 import terraform.hcl
 
@@ -51,6 +51,26 @@ def merge(a: TerraformModule, b: TerraformModule) -> TerraformModule:
 
     return merged
 
+def files_in_module(path: Path) -> Iterable[Path]:
+
+    if os.environ.get('OPENTOFU') == 'true':
+        files = {}
+
+        for filename in path.iterdir():
+            stem = filename.stem
+
+            if filename.suffix == '.tf' and stem not in files:
+                files[stem] = filename
+            elif filename.suffix == '.tofu':
+                files[stem] = filename
+
+        yield from files.values()
+
+    else:
+        for filename in path.iterdir():
+            if filename.suffix == '.tf':
+                yield filename
+
 
 def load_module(path: Path) -> TerraformModule:
     """
@@ -62,12 +82,10 @@ def load_module(path: Path) -> TerraformModule:
 
     module = cast(TerraformModule, {})
 
-    for file in os.listdir(path):
-        if not file.endswith('.tf'):
-            continue
+    for file in files_in_module(path):
 
         try:
-            tf_file = cast(TerraformModule, terraform.hcl.load(os.path.join(path, file)))
+            tf_file = cast(TerraformModule, terraform.hcl.load(file))
             module = merge(module, tf_file)
         except Exception as e:
             # ignore tf files that don't parse

tests/test_module.py

@@ -1,5 +1,7 @@
+from pathlib import Path
+
 from terraform.hcl import loads
-from terraform.module import get_sensitive_variables
+from terraform.module import get_sensitive_variables, files_in_module
 
 
 def test_get_sensitive_variables():
@@ -26,3 +28,17 @@ def test_get_sensitive_variables():
 ''')
 
     assert get_sensitive_variables(module) == ['secret', 'super_secret']
+
+def test_load_terraform_module():
+    assert set(s.name for s in files_in_module(Path('tests/tofu-module'))) == {
+        'blah.tf',
+        'hello.tf',
+    }
+
+def test_load_tofu_module(monkeypatch):
+    monkeypatch.setenv('OPENTOFU', 'true')
+    assert set(s.name for s in files_in_module(Path('tests/tofu-module'))) == {
+        'blah.tf',
+        'hello.tofu',
+        'tofu-only.tofu'
+    }