Minor improvements in diamonds.py example