fix: processing two-character elements (#1067)

deepmodeling · Mar 25, 2021 · 752bfbf · 752bfbf · vercel · Mar 25, 2021
1 parent f142499
commit 752bfbf
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/reacnetgenerator/_reachtml.py b/reacnetgenerator/_reachtml.py
@@ -58,9 +58,11 @@ def _re(self, smi):
         [C]
         >>> self._re('[CH]')
         [CH]
+        >>> self._re('Na')
+        [Na]
         """
-        elements = "".join([an.upper() + an.lower() for an in self.atomname if an != 'H'])
-        smi = re.sub(r'(?<!\[)([' + elements + r'])(?!H)', r'[\1]', smi)
+        elements = "|".join([((an.upper() + "|" + an.lower()) if len(an)==1 else an) for an in self.atomname if an != 'H'])
+        smi = re.sub(r'(?<!\[)(' + elements + r')(?!H)', r'[\1]', smi)
         return smi.replace("[HH]", "[H]")
 
     def _handlereaction(self, line):