Estimating AF from DP4 in a VCF file

I was using a VCF generated using samtools mpileup which had the DP4 column but no AF column. The AF column is the allele fraction which is the fraction of the alternate allele read count, n_alt_allele_reads/n_total_reads

bcftools does have a plugin to calculate this, bcftools +fill-tags file.vcf but it does require the GT field to be present in the VCF. I instead decided to the DP4 field which gives the number of reads supporting the reference and alternate alleles on the forward and reverse strands. Most of the sites that I was looking at are bi-allelic. I used the PyVCF module to parse the vcf file in python, the code is below,

import vcf

vcf_reader = vcf.VCFReader(open('samtools1.1_output.vcf'))
vcf_writer = vcf.Writer(open('/dev/stdout', 'w'), vcf_reader)

for rec in vcf_reader:
    DP4 = rec.samples[0]['DP4']
    ref_total = float(DP4[0] + DP4[1])
    alt_total = float(DP4[2] + DP4[3])
    AF = alt_total/(ref_total + alt_total)
    if AF > 0.01:
        rec.INFO['AF'] = AF
        vcf_writer.write_record(rec)