diff --git a/nodebook/nodebookcore.py b/nodebook/nodebookcore.py
index 98b5cdc..ddae374 100644
--- a/nodebook/nodebookcore.py
+++ b/nodebook/nodebookcore.py
@@ -46,6 +46,9 @@ class ReferenceFinder(ast.NodeVisitor):
         self.locals.add(node.name)
         self.generic_visit(node)
 
+    def visit_arg(self, node):
+        self.locals.add(node.arg)
+
     def visit_AugAssign(self, node):
         target = node.target
         while (type(target) is ast.Subscript):
diff --git a/setup.py b/setup.py
index 11d6a77..adc7229 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@ import sys
 
 setup(
     name='nodebook',
-    version='0.2.0',
+    version='0.2.1',
     author='Kevin Zielnicki',
     author_email='kzielnicki@stitchfix.com',
     license='Stitch Fix 2017',

Functions don't work in nodebook in py3
Because of changes to the ast in Python 3, functions no longer are parsed correctly. Eg, from @hacktuarial:
```
def add(a, b):
    return a + b
```
nodebook throws an error:
```
KeyError: "name 'a' is not defined"
```
diff --git a/tests/test_nodebookcore.py b/tests/test_nodebookcore.py
index cfd9646..470121c 100644
--- a/tests/test_nodebookcore.py
+++ b/tests/test_nodebookcore.py
@@ -42,6 +42,16 @@ class TestReferenceFinder(object):
         assert rf.locals == {'pd', 'y'}
         assert rf.imports == {'pandas'}
 
+    def test_function(self, rf):
+        code_tree = ast.parse(
+            "def add(x,y):\n"
+            "    return x+y\n"
+        )
+        rf.visit(code_tree)
+        assert rf.inputs == set()
+        assert rf.locals == {'add', 'x', 'y'}
+        assert rf.imports == set()
+
 
 class TestNodebook(object):
     @pytest.fixture()
diff --git a/tests/test_pickledict.py b/tests/test_pickledict.py
index ef35fdd..90b7088 100644
--- a/tests/test_pickledict.py
+++ b/tests/test_pickledict.py
@@ -33,6 +33,12 @@ class TestPickleDict(object):
         df = pd.DataFrame({'a': [0, 1, 2], 'b': ['foo', 'bar', 'baz']})
         mydict['test_df'] = df
         assert mydict['test_df'].equals(df)
+        
+    def test_func(self, mydict):
+        def add(a, b):
+            return a + b
+        mydict['test_func'] = add
+        assert mydict['test_func'](3,5) == 8
 
     def test_immutability(self, mydict):
         l = [1, 2, 3]

pip install -e .[dev]
parse_log_pytest
pytest
pytest
apt-get update apt-get install -y gcc
3.9
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
anyio==4.9.0
argon2-cffi==23.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
asttokens==3.0.0
async-lru==2.0.5
attrs==25.3.0
babel==2.17.0
beautifulsoup4==4.13.3
bleach==6.2.0
certifi==2025.1.31
cffi==1.17.1
charset-normalizer==3.4.1
click==8.1.8
comm==0.2.2
debugpy==1.8.13
decorator==5.2.1
defusedxml==0.7.1
dill==0.3.9
exceptiongroup @ file:///croot/exceptiongroup_1706031385326/work
executing==2.2.0
fastjsonschema==2.21.1
fqdn==1.5.1
h11==0.14.0
httpcore==1.0.7
httpx==0.28.1
idna==3.10
importlib_metadata==8.6.1
iniconfig @ file:///home/linux1/recipes/ci/iniconfig_1610983019677/work
ipykernel==6.29.5
ipython==8.18.1
ipywidgets==8.1.5
isoduration==20.11.0
jedi==0.19.2
Jinja2==3.1.6
json5==0.10.0
jsonpointer==3.0.0
jsonschema==4.23.0
jsonschema-specifications==2024.10.1
jupyter==1.1.1
jupyter-console==6.6.3
jupyter-events==0.12.0
jupyter-lsp==2.2.5
jupyter_client==8.6.3
jupyter_core==5.7.2
jupyter_server==2.15.0
jupyter_server_terminals==0.5.3
jupyterlab==4.3.6
jupyterlab_pygments==0.3.0
jupyterlab_server==2.27.3
jupyterlab_widgets==3.0.13
MarkupSafe==3.0.2
matplotlib-inline==0.1.7
mistune==3.1.3
msgpack-python==0.5.6
nbclient==0.10.2
nbconvert==7.16.6
nbformat==5.10.4
nest-asyncio==1.6.0
-e git+https://github.com/stitchfix/nodebook.git@b6e1ec614fd39acb740b04e99ee7e97d99122420#egg=nodebook
notebook==7.3.3
notebook_shim==0.2.4
numpy==2.0.2
overrides==7.7.0
packaging @ file:///croot/packaging_1734472117206/work
pandas==2.2.3
pandocfilters==1.5.1
parso==0.8.4
pexpect==4.9.0
platformdirs==4.3.7
pluggy @ file:///croot/pluggy_1733169602837/work
prometheus_client==0.21.1
prompt_toolkit==3.0.50
psutil==7.0.0
ptyprocess==0.7.0
pure_eval==0.2.3
pycparser==2.22
Pygments==2.19.1
pytest @ file:///croot/pytest_1738938843180/work
pytest-runner==6.0.1
python-dateutil==2.9.0.post0
python-json-logger==3.3.0
pytz==2025.2
PyYAML==6.0.2
pyzmq==26.3.0
referencing==0.36.2
requests==2.32.3
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rpds-py==0.24.0
Send2Trash==1.8.3
six==1.17.0
sniffio==1.3.1
soupsieve==2.6
stack-data==0.6.3
terminado==0.18.1
tinycss2==1.4.0
tomli @ file:///opt/conda/conda-bld/tomli_1657175507142/work
tornado==6.4.2
traitlets==5.14.3
types-python-dateutil==2.9.0.20241206
typing_extensions==4.13.0
tzdata==2025.2
uri-template==1.3.0
urllib3==2.3.0
wcwidth==0.2.13
webcolors==24.11.1
webencodings==0.5.1
websocket-client==1.8.0
widgetsnbextension==4.0.13
zipp==3.21.0

name: nodebook
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - exceptiongroup=1.2.0=py39h06a4308_0
  - iniconfig=1.1.1=pyhd3eb1b0_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=3.0.16=h5eee18b_0
  - packaging=24.2=py39h06a4308_0
  - pip=25.0=py39h06a4308_0
  - pluggy=1.5.0=py39h06a4308_0
  - pytest=8.3.4=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - setuptools=75.8.0=py39h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - tomli=2.0.1=py39h06a4308_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - anyio==4.9.0
      - argon2-cffi==23.1.0
      - argon2-cffi-bindings==21.2.0
      - arrow==1.3.0
      - asttokens==3.0.0
      - async-lru==2.0.5
      - attrs==25.3.0
      - babel==2.17.0
      - beautifulsoup4==4.13.3
      - bleach==6.2.0
      - certifi==2025.1.31
      - cffi==1.17.1
      - charset-normalizer==3.4.1
      - click==8.1.8
      - comm==0.2.2
      - debugpy==1.8.13
      - decorator==5.2.1
      - defusedxml==0.7.1
      - dill==0.3.9
      - executing==2.2.0
      - fastjsonschema==2.21.1
      - fqdn==1.5.1
      - h11==0.14.0
      - httpcore==1.0.7
      - httpx==0.28.1
      - idna==3.10
      - importlib-metadata==8.6.1
      - ipykernel==6.29.5
      - ipython==8.18.1
      - ipywidgets==8.1.5
      - isoduration==20.11.0
      - jedi==0.19.2
      - jinja2==3.1.6
      - json5==0.10.0
      - jsonpointer==3.0.0
      - jsonschema==4.23.0
      - jsonschema-specifications==2024.10.1
      - jupyter==1.1.1
      - jupyter-client==8.6.3
      - jupyter-console==6.6.3
      - jupyter-core==5.7.2
      - jupyter-events==0.12.0
      - jupyter-lsp==2.2.5
      - jupyter-server==2.15.0
      - jupyter-server-terminals==0.5.3
      - jupyterlab==4.3.6
      - jupyterlab-pygments==0.3.0
      - jupyterlab-server==2.27.3
      - jupyterlab-widgets==3.0.13
      - markupsafe==3.0.2
      - matplotlib-inline==0.1.7
      - mistune==3.1.3
      - msgpack-python==0.5.6
      - nbclient==0.10.2
      - nbconvert==7.16.6
      - nbformat==5.10.4
      - nest-asyncio==1.6.0
      - notebook==7.3.3
      - notebook-shim==0.2.4
      - numpy==2.0.2
      - overrides==7.7.0
      - pandas==2.2.3
      - pandocfilters==1.5.1
      - parso==0.8.4
      - pexpect==4.9.0
      - platformdirs==4.3.7
      - prometheus-client==0.21.1
      - prompt-toolkit==3.0.50
      - psutil==7.0.0
      - ptyprocess==0.7.0
      - pure-eval==0.2.3
      - pycparser==2.22
      - pygments==2.19.1
      - pytest-runner==6.0.1
      - python-dateutil==2.9.0.post0
      - python-json-logger==3.3.0
      - pytz==2025.2
      - pyyaml==6.0.2
      - pyzmq==26.3.0
      - referencing==0.36.2
      - requests==2.32.3
      - rfc3339-validator==0.1.4
      - rfc3986-validator==0.1.1
      - rpds-py==0.24.0
      - send2trash==1.8.3
      - six==1.17.0
      - sniffio==1.3.1
      - soupsieve==2.6
      - stack-data==0.6.3
      - terminado==0.18.1
      - tinycss2==1.4.0
      - tornado==6.4.2
      - traitlets==5.14.3
      - types-python-dateutil==2.9.0.20241206
      - typing-extensions==4.13.0
      - tzdata==2025.2
      - uri-template==1.3.0
      - urllib3==2.3.0
      - wcwidth==0.2.13
      - webcolors==24.11.1
      - webencodings==0.5.1
      - websocket-client==1.8.0
      - widgetsnbextension==4.0.13
      - zipp==3.21.0
prefix: /opt/conda/envs/nodebook

tests/test_nodebookcore.py::TestReferenceFinder::test_function
tests/test_nodebookcore.py::TestReferenceFinder::test_assign tests/test_nodebookcore.py::TestReferenceFinder::test_augassign tests/test_nodebookcore.py::TestReferenceFinder::test_import tests/test_nodebookcore.py::TestReferenceFinder::test_multiline tests/test_nodebookcore.py::TestNodebook::test_single_node tests/test_nodebookcore.py::TestNodebook::test_node_chain tests/test_pickledict.py::TestPickleDict::test_int[mode_memory] tests/test_pickledict.py::TestPickleDict::test_int[mode_disk] tests/test_pickledict.py::TestPickleDict::test_string[mode_memory] tests/test_pickledict.py::TestPickleDict::test_string[mode_disk] tests/test_pickledict.py::TestPickleDict::test_bytes[mode_memory] tests/test_pickledict.py::TestPickleDict::test_bytes[mode_disk] tests/test_pickledict.py::TestPickleDict::test_df[mode_memory] tests/test_pickledict.py::TestPickleDict::test_df[mode_disk] tests/test_pickledict.py::TestPickleDict::test_func[mode_memory] tests/test_pickledict.py::TestPickleDict::test_func[mode_disk] tests/test_pickledict.py::TestPickleDict::test_immutability[mode_memory] tests/test_pickledict.py::TestPickleDict::test_immutability[mode_disk]
nodebook/nodebookcore.py setup.py
nodebook/nodebookcore.py setup.py