Parsing molecular formula Is it possible to improve this code in terms of efficiency / line number ? (avoid nested code ...). It comes from a test I failed. I don't know how I could have done it better.

from collections import Counter

def parse_molecule(molecule):

    array = [[]]

    for token in map(str, molecule):
        if token.isalpha() and token.istitle():
            last = [token]
            upper = token
            array[-1].append(token)
        elif token.isalpha():
            last = upper + token
            array[-1] = [last]
        elif token.isdigit():
            array[-1].extend(last*(int(token)-1))
        elif token == '(' or token == '[':
            array.append([])
        elif token == ')' or token == ']':
            last = array.pop()
            array[-1].extend(last)

    return dict(Counter(array[-1]))


Examples of function usage are following:

&gt;&gt;&gt; water = 'H2O'
&gt;&gt;&gt; parse_molecule(water)
&gt;&gt;&gt;{'H': 2, 'O': 1}

&gt;&gt;&gt; magnesium_hydroxide = 'Mg(OH)2'
&gt;&gt;&gt; parse_molecule(magnesium_hydroxide)
&gt;&gt;&gt; {'Mg': 1, 'O': 2, 'H': 2}

&gt;&gt;&gt; fremy_salt = 'K4[ON(SO3)2]2'
&gt;&gt;&gt; parse_molecule(fremy_salt)
&gt;&gt;&gt; {'K': 4, 'O': 14, 'N': 2, 'S': 4}


This is the statement:  


  For a given chemical formula represented by a string, count the number
  of atoms of each element contained in the molecule and return a dict.


For example:


water = 'H2O'
parse_molecule(water)
return {'H': 2, 'O': 1}

magnesium_hydroxide = 'Mg(OH)2'&lt;br&gt;
parse_molecule(magnesium_hydroxide)&lt;br&gt;
return {'Mg': 1, 'O': 2, 'H': 2}

fremy_salt = 'K4[ON(SO3)2]2'&lt;br&gt;
parse_molecule(fremy_salt)&lt;br&gt;
return {'K': 4, 'O': 14, 'N': 2, 'S': 4}



As you can see, some formulas have brackets in them. The index outside the brackets tells you that you have to multiply count of each atom inside the bracket on this index. For example, in Fe(NO3)2 you have one iron atom, two nitrogen atoms and six oxygen atoms.

Note that brackets may be round, square or curly and can also be nested. Index after the braces is optional.
This code is confusing and inefficient:

for token in map(str, molecule):


It treats molecule as some kind of iterable object, and applies the str function to each element the iteration returns, an assigns each result successively to token.  In this case, molecule is a string, which is already a sequence of iterable characters, and a single character is itself a string, so this map function is taking individual character strings and calling str() on those character strings to return the same character strings.

The statement could be written much simpler:

for token in molecule:




These statements are testing for only 2 of the three grouping brackets:

    elif token == '(' or token == '[':
        ...
    elif token == ')' or token == ']':
        ...


The correct statements would read:

    elif token == '(' or token == '[' or token == '{':
        ...
    elif token == ')' or token == ']' or token == '}':
        ...


however, this is rather verbose.  You can use the in operator to check if a character appears anywhere in a string, so these test could be written more compactly as:

    elif token in '([{':
        ...
    elif token in ')]}':
        ...




isdigit() is not the best function to use to verify a string can be parsed as an integer.  It includes superscript and subscript digits, and other characters which will cause an exception when passed into int(...):

&gt;&gt;&gt; "³₂".isdigit()
True


Instead, use isdecimal().



As mention by scnerd in their answer, it is hard to parse strings character by character.  You are detecting a lower case letter, adding it to the last token you found, and replacing the last token.  You would need to do the same thing when you find a digit, and only when the current series of digits end perform the replication of the last radical, but this results in numerous special cases, which again make the parsing hard.

It is better to break the molecule into low level tokens, not individual characters, and process complete tokens.  As mention by scnerd, the regular expression engine can help.  They diverted into attempt to match parenthesis in the regular expression parsing, which complicates things far too much.  Here is a simpler approach:

import re

def tokenize_molecule(molecule):
    return re.findall('[A-Z][a-z]?|\d+|.', molecule)

molecules = ['H2O', 'Mg(OH)2', 'K4[ON(SO3)2]2', 'C12H22O11']

for molecule in molecules:
    print(molecule)
    print("\t", tokenize_molecule(molecule))


which outputs:

H2O
     ['H', '2', 'O']
Mg(OH)2
     ['Mg', '(', 'O', 'H', ')', '2']
K4[ON(SO3)2]2
     ['K', '4', '[', 'O', 'N', '(', 'S', 'O', '3', ')', '2', ']', '2']
C12H22O11
     ['C', '12', 'H', '22', 'O', '11']


Here, the regex has 3 parts: [A-Z][a-z]? which matches an element, \d+, which matches one or more digits, and . which matches any character.  These are joined together with the | alternative operator; if the current character does not begin an element name, or a digit group, it is returned as a single character ... which is one of the bracket characters in a properly formed molecule.

Using this initial parsing will make the parsing the molecule much easier:

def parse_molecule(molecule):

    # ... 

    for token in re.findall('[A-Z][a-z]?|\d+|.', molecule):

        if token.isalpha():
            # ...
        elif token.isdecimal():
            count = int(token)
            # ...
        elif token in '([{':
            # ...
        elif token in ')]}':
            # ...
        else:
            raise ValueError('Unrecognized character in molecule')

    # ...