Add extra functionality based on #20 by @kidpollo

Tests included this time.