From ac042d78d733633d6944a92fcf63cd6441894258 Mon Sep 17 00:00:00 2001
From: chhzh123 <hc676@cornell.edu>
Date: Wed, 11 Dec 2024 22:11:18 -0500
Subject: [PATCH] Fix

---
 allo/customize.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/allo/customize.py b/allo/customize.py
index 5e33b260..98ebd94d 100644
--- a/allo/customize.py
+++ b/allo/customize.py
@@ -73,10 +73,6 @@ def getsourcefile(obj):
     return ret
 
 
-def getsourcelines(obj):
-    return inspect.getsourcelines(obj)
-
-
 def wrapped_apply(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
@@ -903,9 +899,9 @@ def customize(
 ):
     # Get Python AST
     if isinstance(fn, str):
-        src = fn
+        src, starting_line_no = fn, 1
     else:
-        src, starting_line_no = getsourcelines(fn)
+        src, starting_line_no = inspect.getsourcelines(fn)
         src = [textwrap.fill(line, tabsize=4, width=9999) for line in src]
         src = textwrap.dedent("\n".join(src))
     tree = parse_ast(src, starting_line_no=starting_line_no, verbose=verbose)