cginclude watcher script for editing unity shaders

cgincWatcher.py

#!/usr/bin/env python

# CGInclude Watcher
# Joel Davis (joeld42@gmail.com, @joeld42 on Twitter)
# 
# Usage: Run this from a directory where you have shaders.
#
# Use at your own risk (and use an editor that properly handles
# open files being modified by another program). 
#
# This script scans a directory of Unity .shader files, and 
# when shaders include another file in the watched directory
# (typically CGIncludes) it adds a hash to the include line. For 
# example:
#
#  #include "CGIncludes/CommonEffect.cginc"   // A64C87A
#
# Whenever the contents of that include file change, the script will
# update the .shaders that use it and modify the hash. This will in
# turn trigger Unity to know it should rebuild the shader even if
# only the include changed.

import os, sys
import re
import hashlib
import time

# This is the name of the path where the include files we are
# editing live. Only includes in this directory are watched.
WATCH_INCLUDE_PATH = "CGIncludes"



reIncludeStuff = re.compile( r".*\#include\s+\"(?P<incfile>.*)\"(.*//(?P<hash>[A-z0-9]+))?.*")

# Not used, but you might want to use this if you're using something that
# cares about
def fileNeedsUpdate( shaderFile, includeFile ):

    incmtime = os.path.getmtime( includeFile )
    srcmtime = os.path.getmtime( shaderFile )

    if incmtime < srcmtime:
        return True

    # file is up to date
    return False

# computes a really good hash and then truncates it to a bad one
def sha256sum(filename):
    h = hashlib.sha256()
    with open(filename, 'rb', buffering=0) as f:
        for b in iter(lambda : f.read(128*1024), b''):
            h.update(b)

    return h.hexdigest()[:8] # Use truncated hash 

# Rewrites a shader file, updating the hashes on the includes
def updateIncludeHash( filename, cgincHash ):
	fn2 = "/var/tmp/tmp_" + filename
	print "Rewriting ", filename, fn2 

	fp = open( fn2, "wt" )
	for line in open(filename).readlines():
		m = reIncludeStuff.match( line )
		if m:
			incfile = m.group("incfile") 
			if incfile and cgincHash.has_key( incfile ):
				suffixPos = line.find( '.cginc' )
				if (suffixPos == -1):
					print "Can't find suffix?"
					return

				line = line[:suffixPos+7]

				line = line + " //" + cgincHash[incfile]+"\n"

		fp.write(line)
	fp.close()

	# Atomic replace
	os.rename( fn2, filename )



if __name__=='__main__':

	shaderIncludes = {}	
	cgincHash = {}

	# find .shader files
	for f in os.listdir('.'):
		fn, ext = os.path.splitext(f)
		if ext =='.shader':
			# scan shader and look for includes
			for line in open(f).readlines():
				m = reIncludeStuff.match( line )
				if (m):					
					includeFile = m.group("incfile")
					hashVal = m.group("hash")
					if hashVal == None:
						hashVal = "ABCDE" # No stored hash, use a dummy value

					# Only care about stuff in CGIncludes
					if includeFile.startswith( WATCH_INCLUDE_PATH ):
						#print includeFile
						if not shaderIncludes.has_key( f ):
							shaderIncludes[f] = []

							cgincHash[ includeFile ] = hashVal

						shaderIncludes[f].append( includeFile)


	shaderFiles = shaderIncludes.keys()
	shaderFiles.sort()

	for f in shaderFiles:
		print f, shaderIncludes[f]

	print cgincHash

	# main loop, check for updates

	while 1:
		time.sleep(3)
		needsUpdate = set()
		for inc in cgincHash.keys():		
			newhash = sha256sum( inc )
			if newhash != cgincHash[inc]:
				cgincHash[inc] = newhash

				for f in shaderFiles:
					if inc in shaderIncludes[f]:
						needsUpdate.add( f )

		if len(needsUpdate):

			for fn in list(needsUpdate):
				updateIncludeHash( fn, cgincHash )

			print "up to date..."
			#sys.exit(1)