diff --git a/BSD3.txt b/BSD3.txt
--- a/BSD3.txt
+++ b/BSD3.txt
@@ -1,5 +1,5 @@
 Copyright (c) 2010, Oscar Finnsson
-Copyright (c) 2013-2014, Benno Fünfstück
+Copyright (c) 2013-2017, Benno Fünfstück
 All rights reserved.
 
 Redistribution and use in source and binary forms, with or without
diff --git a/example-explicit.hs b/example-explicit.hs
new file mode 100644
--- /dev/null
+++ b/example-explicit.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE TemplateHaskell #-}
+module Main where
+
+import Test.Tasty
+import Test.Tasty.TH
+import Test.Tasty.QuickCheck
+import Test.Tasty.HUnit
+
+main :: IO ()
+main = $(defaultMainGeneratorFor "explicit" ["prop_length_append", "case_length_1", "test_plus"])
+
+prop_length_append :: [Int] -> [Int] -> Bool
+prop_length_append as bs = length (as ++ bs) == length as + length bs
+
+case_length_1 :: Assertion
+case_length_1 = 1 @=? length [()]
+
+case_add :: Assertion
+case_add = 7 @=? (3 + 4)
+
+test_plus :: [TestTree]
+test_plus =
+  [ $(testGroupGeneratorFor "case_add" ["case_add"])
+    -- ...
+  ]
diff --git a/example-literate.lhs b/example-literate.lhs
new file mode 100644
--- /dev/null
+++ b/example-literate.lhs
@@ -0,0 +1,36 @@
+This is an example of using tasty-th with literate haskell files.
+First, we need to import the library and enable template haskell:
+
+> {-# LANGUAGE TemplateHaskell #-}
+> import Test.Tasty
+> import Test.Tasty.TH
+> import Test.Tasty.QuickCheck
+> import Test.Tasty.HUnit
+
+Now, we can write some quickcheck properties:
+
+> prop_length_append :: [Int] -> [Int] -> Bool
+> prop_length_append as bs = length (as ++ bs) == length as + length bs
+
+Or write a HUnit test case:
+
+> case_length_1 :: Assertion
+> case_length_1 = 1 @=? length [()]
+
+Properties in comments are not run:
+
+prop_comment :: Assertion
+prop_comment = assertFailure "property in comment should not be run"
+
+We can also create test trees:
+
+> test_plus :: [TestTree]
+> test_plus =
+>   [ testCase "3 + 4" (7 @=? (3 + 4))
+>     --  ...
+>   ]
+
+We only need a main now that collects all our tests:
+
+> main :: IO ()
+> main = $(defaultMainGenerator)
diff --git a/example.hs b/example.hs
new file mode 100644
--- /dev/null
+++ b/example.hs
@@ -0,0 +1,29 @@
+{-# LANGUAGE TemplateHaskell #-}
+module Main where
+
+import Test.Tasty
+import Test.Tasty.TH
+import Test.Tasty.QuickCheck
+import Test.Tasty.HUnit
+
+main :: IO ()
+main = $(defaultMainGenerator)
+
+{-
+Properties in comments are not run:
+
+prop_comment :: Assertion
+prop_comment = assertFailure "property in comment should not be run"
+-}
+
+prop_length_append :: [Int] -> [Int] -> Bool
+prop_length_append as bs = length (as ++ bs) == length as + length bs
+
+case_length_1 :: Assertion
+case_length_1 = 1 @=? length [()]
+
+test_plus :: [TestTree]
+test_plus =
+  [ testCase "3 + 4" (7 @=? (3 + 4))
+    -- ...
+  ]
diff --git a/src/Test/Tasty/TH.hs b/src/Test/Tasty/TH.hs
--- a/src/Test/Tasty/TH.hs
+++ b/src/Test/Tasty/TH.hs
@@ -12,103 +12,141 @@
 -----------------------------------------------------------------------------
 {-# LANGUAGE TemplateHaskell #-}
 
+-- | This module provides TemplateHaskell functions to automatically generate
+-- tasty TestTrees from specially named functions. See the README of the package
+-- for examples.
+--
+-- Important: due to to the GHC staging restriction, you must put any uses of these
+-- functions at the end of the file, or you may get errors due to missing definitions.
 module Test.Tasty.TH
   ( testGroupGenerator
   , defaultMainGenerator
+  , testGroupGeneratorFor
+  , defaultMainGeneratorFor
+  , extractTestFunctions
+  , locationModule
   ) where
 
+import Control.Monad (join)
+import Control.Applicative
+import Language.Haskell.Exts (parseFileContentsWithMode)
+import Language.Haskell.Exts.Parser (ParseResult(..), defaultParseMode, parseFilename)
+import qualified Language.Haskell.Exts.Syntax as S
 import Language.Haskell.TH
-import Language.Haskell.Extract
+import Data.Maybe
+import Data.Data (gmapQ, Data)
+import Data.Typeable (cast)
+import Data.List (nub, isPrefixOf, find)
+import qualified Data.Foldable as F
 
 import Test.Tasty
+import Prelude
 
--- | Generate the usual code and extract the usual functions needed in order to run HUnit/Quickcheck/Quickcheck2.
--- All functions beginning with case_, prop_ or test_ will be extracted.
+-- | Convenience function that directly generates an `IO` action that may be used as the
+-- main function. It's just a wrapper that applies 'defaultMain' to the 'TestTree' generated
+-- by 'testGroupGenerator'.
 --
--- > {-# OPTIONS_GHC -fglasgow-exts -XTemplateHaskell #-}
--- > module MyModuleTest where
--- > import Test.HUnit
--- > import MainTestGenerator
--- >
--- > main = $(defaultMainGenerator)
--- >
--- > case_Foo = do 4 @=? 4
--- >
--- > case_Bar = do "hej" @=? "hej"
--- >
--- > prop_Reverse xs = reverse (reverse xs) == xs
--- > where types = xs :: [Int]
--- >
--- > test_Group =
--- > [ testCase "1" case_Foo
--- > , testProperty "2" prop_Reverse
--- > ]
+-- Example usage:
 --
--- will automagically extract prop_Reverse, case_Foo, case_Bar and test_Group and run them as well as present them as belonging to the testGroup 'MyModuleTest' such as
+-- @
+-- -- properties, test cases, ....
 --
--- > me: runghc MyModuleTest.hs
--- > MyModuleTest:
--- > Reverse: [OK, passed 100 tests]
--- > Foo: [OK]
--- > Bar: [OK]
--- > Group:
--- > 1: [OK]
--- > 2: [OK, passed 100 tests]
--- >
--- > Properties Test Cases Total
--- > Passed 2 3 5
--- > Failed 0 0 0
--- > Total 2 3 5
+-- main :: IO ()
+-- main = $('defaultMainGenerator')
+-- @
 defaultMainGenerator :: ExpQ
-defaultMainGenerator =
-  [| defaultMain $ testGroup $(locationModule) $ $(propListGenerator) ++ $(caseListGenerator) ++ $(testListGenerator)|]
+defaultMainGenerator = [| defaultMain $(testGroupGenerator) |]
 
--- | Generate the usual code and extract the usual functions needed for a testGroup in HUnit/Quickcheck/Quickcheck2.
---   All functions beginning with case_, prop_ or test_ will be extracted.
+-- | This function generates a 'TestTree' from functions in the current module. 
+-- The test tree is named after the current module.
 --
---   > -- file SomeModule.hs
---   > fooTestGroup = $(testGroupGenerator)
---   > main = defaultMain [fooTestGroup]
---   > case_1 = do 1 @=? 1
---   > case_2 = do 2 @=? 2
---   > prop_p xs = reverse (reverse xs) == xs
---   >  where types = xs :: [Int]
+-- The following definitions are collected by `testGroupGenerator`:
 --
---   is the same as
+-- * a test_something definition in the current module creates a sub-testGroup with the name "something"
+-- * a prop_something definition in the current module is added as a QuickCheck property named "something"
+-- * a case_something definition leads to a HUnit-Assertion test with the name "something"
 --
---   > -- file SomeModule.hs
---   > fooTestGroup = testGroup "SomeModule" [testProperty "p" prop_1, testCase "1" case_1, testCase "2" case_2]
---   > main = defaultMain [fooTestGroup]
---   > case_1 = do 1 @=? 1
---   > case_2 = do 2 @=? 2
---   > prop_1 xs = reverse (reverse xs) == xs
---   >  where types = xs :: [Int]
+-- Example usage:
+--
+-- @
+-- prop_example :: Int -> Int -> Bool
+-- prop_example a b = a + b == b + a
+--
+-- tests :: 'TestTree'
+-- tests = $('testGroupGenerator')
+-- @
 testGroupGenerator :: ExpQ
-testGroupGenerator =
-  [| testGroup $(locationModule) $ $(propListGenerator) ++ $(caseListGenerator) ++ $(testListGenerator) |]
+testGroupGenerator = join $ testGroupGeneratorFor <$> fmap loc_module location <*> testFunctions
+ where
+  testFunctions = location >>= runIO . extractTestFunctions . loc_filename
 
-listGenerator :: String -> String -> ExpQ
-listGenerator beginning funcName =
-  functionExtractorMap beginning (applyNameFix funcName)
+-- | Retrieves all function names from the given file that would be discovered by 'testGroupGenerator'.
+extractTestFunctions :: FilePath -> IO [String]
+extractTestFunctions filePath = do
+  file <- readFile filePath
+  -- we first try to parse the file using haskell-src-exts
+  -- if that fails, we fallback to lexing each line, which is less
+  -- accurate but is more reliable (haskell-src-exts sometimes struggles
+  -- with less-common GHC extensions).
+  let functions = fromMaybe (lexed file) (parsed file)
+      filtered pat = filter (pat `isPrefixOf`) functions
+  return . nub $ concat [filtered "prop_", filtered "case_", filtered "test_"]
+ where
+  lexed = map fst . concatMap lex . lines
+  
+  parsed file = case parseFileContentsWithMode (defaultParseMode { parseFilename = filePath }) file of
+    ParseOk parsedModule -> Just (declarations parsedModule)
+    ParseFailed _ _ -> Nothing
+  declarations (S.Module _ _ _ _ decls) = concatMap testFunName decls
+  declarations _ = []
+  testFunName (S.PatBind _ pat _ _) = patternVariables pat
+  testFunName (S.FunBind _ clauses) = nub (map clauseName clauses)
+  testFunName _ = []
+  clauseName (S.Match _ name _ _ _) = nameString name
+  clauseName (S.InfixMatch _ _ name _ _ _) = nameString name
 
-propListGenerator :: ExpQ
-propListGenerator = listGenerator "^prop_" "testProperty"
+-- | Convert a 'Name' to a 'String'
+nameString :: S.Name l -> String
+nameString (S.Ident _ n) = n
+nameString (S.Symbol _ n) = n
 
-caseListGenerator :: ExpQ
-caseListGenerator = listGenerator "^case_" "testCase"
+-- | Find all variables that are bound in the given pattern.
+patternVariables :: Data l => S.Pat l -> [String]
+patternVariables = go
+ where
+  go (S.PVar _ name) = [nameString name]
+  go pat = concat $ gmapQ (F.foldMap go . cast) pat
 
-testListGenerator :: ExpQ
-testListGenerator = listGenerator "^test_" "testGroup"
+-- | Extract the name of the current module.
+locationModule :: ExpQ
+locationModule = do
+  loc <- location
+  return $ LitE $ StringL $ loc_module loc
 
--- | The same as
---   e.g. \n f -> testProperty (fixName n) f
-applyNameFix :: String -> ExpQ
-applyNameFix n =
-  do fn <- [|fixName|]
-     return $ LamE [VarP (mkName "n")] (AppE (VarE (mkName n)) (AppE fn (VarE (mkName "n"))))
+-- | Like 'testGroupGenerator', but generates a test group only including the specified function names.
+-- The function names still need to follow the pattern of starting with one of @prop_@, @case_@ or @test_@.
+testGroupGeneratorFor
+  :: String   -- ^ The name of the test group itself
+  -> [String] -- ^ The names of the functions which should be included in the test group
+  -> ExpQ
+testGroupGeneratorFor name functionNames = [| testGroup name $(listE (mapMaybe test functionNames)) |]
+ where
+  testFunctions = [("prop_", "testProperty"), ("case_", "testCase"), ("test_", "testGroup")]
+  getTestFunction fname = snd <$> find ((`isPrefixOf` fname) . fst) testFunctions
+  test fname = do
+    fn <- getTestFunction fname
+    return $ appE (appE (varE $ mkName fn) (stringE (fixName fname))) (varE (mkName fname))
 
+-- | Like 'defaultMainGenerator', but only includes the specific function names in the test group.
+-- The function names still need to follow the pattern of starting with one of @prop_@, @case_@ or @test_@.
+defaultMainGeneratorFor
+  :: String   -- ^ The name of the top-level test group
+  -> [String] -- ^ The names of the functions which should be included in the test group
+  -> ExpQ
+defaultMainGeneratorFor name fns = [| defaultMain $(testGroupGeneratorFor name fns) |]
+
 fixName :: String -> String
-fixName name = replace '_' ' ' $ drop 5 name
+fixName = replace '_' ' ' . tail . dropWhile (/= '_')
 
 replace :: Eq a => a -> a -> [a] -> [a]
 replace b v = map (\i -> if b == i then v else i)
diff --git a/tasty-th.cabal b/tasty-th.cabal
--- a/tasty-th.cabal
+++ b/tasty-th.cabal
@@ -1,22 +1,38 @@
 name: tasty-th
-version: 0.1.2
-cabal-version: >= 1.6
+version: 0.1.7
+cabal-version: >= 1.8
 build-type: Simple
 license: BSD3
 license-file: BSD3.txt
 maintainer: Benno Fünfstück <benno.fuenfstueck@gmail.com>
 homepage: http://github.com/bennofs/tasty-th
-synopsis: Automagically generate the HUnit- and Quickcheck-bulk-code using Template Haskell.
-description: A fork of of test-framework-th modified to use tasty instead of test-framework.
+synopsis: Automatic tasty test case discovery using TH
+description: Generate tasty TestTrees automatically with TemplateHaskell. See the README for example usage.
 
 category: Testing
 author: Oscar Finnsson & Emil Nordling & Benno Fünfstück
+extra-source-files:
+  example.hs
+  example-explicit.hs
+  example-literate.lhs
 
 library
   exposed-modules: Test.Tasty.TH
-  build-depends: base >= 4 && < 5, tasty, language-haskell-extract >= 0.2, template-haskell
+  build-depends: base >= 4 && < 5, haskell-src-exts >= 1.18.0, tasty, template-haskell
   hs-source-dirs: src
+  ghc-options: -Wall
+  other-extensions: TemplateHaskell
 
+test-suite tasty-th-tests
+  hs-source-dirs: tests
+  main-is: Main.hs
+  build-depends:
+    base >= 4 && < 5,
+    tasty-hunit,
+    tasty-th
+  ghc-options: -Wall
+  default-language:    Haskell2010
+  type: exitcode-stdio-1.0
 
 source-repository head
   type:     git
diff --git a/tests/Main.hs b/tests/Main.hs
new file mode 100644
--- /dev/null
+++ b/tests/Main.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE TemplateHaskell #-}
+import Test.Tasty.TH
+import Test.Tasty.HUnit
+import Data.List (sort)
+
+main :: IO ()
+main = $(defaultMainGenerator)
+
+case_example_test_functions :: Assertion
+case_example_test_functions = do
+  functions <- extractTestFunctions "example.hs"
+  let expected = [ "prop_length_append", "case_length_1", "test_plus" ]
+  sort expected @=? sort functions
+
+case_example_explicit_test_functions :: Assertion
+case_example_explicit_test_functions = do
+  functions <- extractTestFunctions "example-explicit.hs"
+  let expected = [ "case_add", "prop_length_append", "case_length_1", "test_plus" ]
+  sort expected @=? sort functions
+
+case_example_literate_test_functions :: Assertion
+case_example_literate_test_functions = do
+  functions <- extractTestFunctions "example-literate.lhs"
+  let expected = [ "prop_length_append", "case_length_1", "test_plus" ]
+  sort expected @=? sort functions