reserve comments

d2lbook/library.py

@@ -8,6 +8,7 @@
 import pathlib
 import ast
 import astor
+from yapf.yapflib.yapf_api import FormatCode
 
 def _write_header(f):
     f.write('# This file is generated automatically through:\n')
@@ -153,17 +154,30 @@ def save_alias(tab_lib):
             f.write('# Alias defined in config.ini\n')
             f.write(alias+'\n\n')
 
-class _FluentAlias(ast.NodeTransformer):
-    def __init__(self, mapping):
-        self._map = {a:b for a, b in mapping}
-
-    def visit_Call(self, node):
-        if isinstance(node.func, ast.Attribute) and isinstance(node.func.value, ast.Name) and node.func.value.id == 'd2l' and node.func.attr in self._map:
-            new_node = ast.Call(
-                ast.Attribute(value=node.args[0], attr=self._map[node.func.attr]),
-                node.args[1:], node.keywords)
-            return self.generic_visit(new_node)
-        return self.generic_visit(node)
+def replace_fluent_alias(source, fluent_mapping):
+    fluent_mapping = {a:b for a, b in fluent_mapping}
+    new_src = source
+    for _ in range(100):  # 100 is a (random) big enough number
+        replaced = False
+        tree = ast.parse(new_src)
+        for node in ast.walk(tree):
+            if (isinstance(node, ast.Call) and
+                isinstance(node.func, ast.Attribute) and
+                isinstance(node.func.value, ast.Name) and
+                node.func.value.id == 'd2l' and
+                node.func.attr in fluent_mapping):
+                new_node = ast.Call(
+                    ast.Attribute(value=node.args[0],
+                                  attr=fluent_mapping[node.func.attr]),
+                    node.args[1:], node.keywords)
+                new_src = new_src.replace(
+                    ast.get_source_segment(new_src, node),
+                    astor.to_source(new_node).rstrip())
+                replaced = True
+                break
+        if not replaced:
+            break
+    return FormatCode(new_src)[0].rstrip()
 
 def replace_alias(nb, tab_lib):
     nb = copy.deepcopy(nb)
@@ -186,8 +200,6 @@ def replace_alias(nb, tab_lib):
             if fluent_mapping:
                 for a, _ in fluent_mapping:
                     if 'd2l.'+a in cell.source:
-                        tree = ast.parse(cell.source)
-                        cell.source = astor.to_source(
-                            _FluentAlias(fluent_mapping).visit(tree)).rstrip()
+                        cell.source = replace_fluent_alias(cell.source, fluent_mapping)
                         break
     return nb

d2lbook/library_test.py

@@ -40,16 +40,18 @@ def test_replace_alias(self):
             ('float(d2l.reduce_sum(d2l.astype(cmp, y.dtype)))',
              'float(cmp.type(y.dtype).sum())'),
             ('\nenc_attention_weights = d2l.reshape(\n    d2l.concat(net.encoder.attention_weights, 0),\n    (num_layers, num_heads, -1, num_steps))\nenc_attention_weights.shape = 2\n',
-             'enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((\n    num_layers, num_heads, -1, num_steps))\nenc_attention_weights.shape = 2'),
+             'enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape(\n    (num_layers, num_heads, -1, num_steps))\nenc_attention_weights.shape = 2'),
             # TODO(mli), a bunch of other cases
             ('float(d2l.reduce_sum(d2l.abs(Y1 - Y2))) < 1e-6',
-             'float(torch.abs(Y1 - Y2).sum()) < 1e-06'),
+             'float(torch.abs(Y1 - Y2).sum()) < 1e-6'),
             ('d2l.plt.scatter(d2l.numpy(features[:, 1]), d2l.numpy(labels), 1);',
-             'd2l.plt.scatter(features[:, (1)].detach().numpy(), labels.detach().numpy(), 1)'),
+             'd2l.plt.scatter(features[:, (1)].detach().numpy(),\n                labels.detach().numpy(), 1)'),
             ('d2l.reshape(multistep_preds[i - tau: i], (1, -1))',
              'multistep_preds[i - tau:i].reshape((1, -1))'),
             ('X = d2l.reshape(d2l.arange(16, dtype=d2l.float32), (1, 1, 4, 4))',
             'X = torch.arange(16, dtype=torch.float32).reshape((1, 1, 4, 4))'),
+            ('# comments\nX = d2l.reshape(a)', '# comments\nX = a.reshape()'),
+            ('X = d2l.reshape(a)  # comments', 'X = a.reshape()  # comments'),
         ]
         for a, b in pairs:
             self.nb.cells[0].source = a

setup.py

@@ -18,13 +18,14 @@
     'gitpython',
     'sphinx_autodoc_typehints',
     'astor',
+    'yapf'
 ]
 
 setup(
     name='d2lbook',
     version=__version__,
     install_requires=requirements,
-    python_requires='>=3.5',
+    python_requires='>=3.8',
     author='D2L Developers',
     author_email='d2l.devs@gmail.com',
     url='https://book.d2l.ai',