How to count Python virtual instructions

Recently, I've had a need to count how many virtual instructions will be executed by the Python virtual machine for a given function or method. It turns out, there is no standard way to do this. But, Python being Python, there is always an easy way out.

The built-in dis module allows us to disassemble the byte-code for any given Python object. However, the challenges are that the results are printed rather than returned. Also, even if the output were returned, we still need to perform some action that will count the number of instructions.

Here is a simple example of what I came up with.

#Virtual instruction count

import sys
import dis

def my_function(op1, op2):
   result=op1**op2
   return result

class VICount:
   def __init__(self, obj):
       self.instructions=[]
       sys.stdout=self
       dis.dis(obj)
       sys.stdout=sys.__stdout__
      
   def write(self, string):
       if string.strip() in dis.opname:
           self.instructions.append(string)
      
   def count(self):
       return len(self.instructions)

print VICount(my_function).count()

In this example, I have a function called my_function(). I would like to determine the number of virtual instructions the Python virtual machine will execute when this function is invoked. For this purpose, I've created a VICount class that is used to count the virtual instructions. The constructor of this class will accept an object to disassemble. We also initialize the list that will store the names of the virtual instructions that are found.

Next, in the constructor, we need to change where the print statements executed by the dis.dis() function go. We do this by changing the sys.stdout file object to self. This is legal since VICount provides a writable file object interface by implementing a write method.

Finally, in the constructor, we need to restore the sys.stdout file object to its' original state.

The write() method is invoked by the print statements executed by dis.dis(). If the string being written is a valid instruction, it is appended to the instruction list.

The count() method simply returns the length of the instruction list. In the case of my_function(), we have six virtual instructions.