Fix IPython import warnings and notebook detection (#86)

- Add type ignore comments for IPython imports
- Fix return type annotation (remove unnecessary quotes)
- Add _is_jupyter() to properly detect notebook environments
- Replace lambda with def function for pylint compliance

Fixes #65

.github/scripts/add-new-checks.sh

@@ -29,11 +29,11 @@ echo "✓ New checks added!"
 echo ""
 echo "Updated required status checks will include:"
 echo "- test (3.10)                    [existing]"
-echo "- test (3.11)                    [existing]" 
+echo "- test (3.11)                    [existing]"
 echo "- test (3.12)                    [existing]"
 echo "- Validate PR Template           [existing]"
 echo "- live-api-tests                 [existing]"
 echo "- ollama-integration-test        [existing]"
 echo "- enforce                        [NEW - linked issue validation]"
 echo "- size                           [NEW - PR size limit]"
-echo "- protect-infrastructure         [NEW - infrastructure file protection]"
+echo "- protect-infrastructure         [NEW - infrastructure file protection]"

.github/scripts/add-size-labels.sh

@@ -23,7 +23,7 @@ gh pr list --limit 50 --json number,additions,deletions --jq '.[]' | while read
     additions=$(echo "$pr_data" | jq -r '.additions')
     deletions=$(echo "$pr_data" | jq -r '.deletions')
     total_changes=$((additions + deletions))
-    
+
     # Determine size label
     if [ $total_changes -lt 50 ]; then
         size_label="size/XS"
@@ -36,20 +36,20 @@ gh pr list --limit 50 --json number,additions,deletions --jq '.[]' | while read
     else
         size_label="size/XL"
     fi
-    
+
     echo "PR #$pr_number: $total_changes lines -> $size_label"
-    
+
     # Remove any existing size labels first
     existing_labels=$(gh pr view $pr_number --json labels --jq '.labels[].name' | grep "^size/" || true)
     if [ ! -z "$existing_labels" ]; then
         echo "  Removing existing label: $existing_labels"
         gh pr edit $pr_number --remove-label "$existing_labels"
     fi
-    
+
     # Add the new size label
     gh pr edit $pr_number --add-label "$size_label"
-    
+
     sleep 1  # Avoid rate limiting
 done
 
-echo "Done adding size labels!"
+echo "Done adding size labels!"

.github/scripts/revalidate-all-prs.sh

@@ -28,7 +28,7 @@ for pr in $PR_NUMBERS; do
     COUNT=$((COUNT + 1))
     echo "[$COUNT/$TOTAL] Triggering revalidation for PR #$pr..."
     gh workflow run revalidate-pr.yml -f pr_number=$pr
-    
+
     # Small delay to avoid rate limiting
     sleep 2
 done
@@ -39,4 +39,4 @@ echo ""
 echo "To monitor progress:"
 echo "  gh run list --workflow=revalidate-pr.yml --limit=$TOTAL"
 echo ""
-echo "To see results, check comments on each PR"
+echo "To see results, check comments on each PR"

.github/workflows/check-linked-issue.yml

@@ -46,9 +46,9 @@ jobs:
               repo: context.repo.repo,
               username: prAuthor
             });
-            
+
             const isMaintainer = ['admin', 'maintain'].includes(authorPermission.permission);
-            
+
             const body = context.payload.pull_request.body || '';
             const match = body.match(/(?:Fixes|Closes|Resolves)\s+#(\d+)/i);
 
@@ -87,4 +87,4 @@ jobs:
               core.setFailed(`Issue #${issueNumber} needs at least ${REQUIRED_THUMBS_UP} 👍 reactions (currently has ${thumbsUp})`);
             } else if (isMaintainer && thumbsUp < REQUIRED_THUMBS_UP) {
               core.info(`Maintainer ${prAuthor} bypassing community support requirement (issue has ${thumbsUp} 👍 reactions)`);
-            }
+            }

.github/workflows/ci.yaml

@@ -154,7 +154,7 @@ jobs:
               repo: context.repo.repo,
               username: context.actor
             });
-            
+
             const isMaintainer = ['admin', 'maintain'].includes(permission.permission);
             if (!isMaintainer) {
               throw new Error(`User ${context.actor} does not have maintainer permissions.`);
@@ -170,7 +170,7 @@ jobs:
         run: |
           git config user.name "github-actions[bot]"
           git config user.email "github-actions[bot]@users.noreply.github.com"
-          
+
           # pull_request_target runs in base repo context, so this is safe
           git fetch origin pull/${{ github.event.pull_request.number }}/head:pr-to-test
           git merge pr-to-test --no-ff --no-edit

.github/workflows/revalidate-pr.yml

@@ -29,14 +29,14 @@ jobs:
               repo: context.repo.repo,
               pull_number: ${{ inputs.pr_number }}
             });
-            
+
             core.info(`Validating PR #${pr.number}: ${pr.title}`);
             core.info(`Author: ${pr.user.login}`);
             core.info(`Changes: +${pr.additions} -${pr.deletions}`);
-            
+
             // Store head SHA for creating status
             core.setOutput('head_sha', pr.head.sha);
-            
+
             return pr;
 
       - name: Create pending status
@@ -60,44 +60,44 @@ jobs:
             const pr = ${{ steps.pr_data.outputs.result }};
             const errors = [];
             let passed = true;
-            
+
             // Check size
             const totalChanges = pr.additions + pr.deletions;
             const MAX_LINES = 1000;
             if (totalChanges > MAX_LINES) {
               errors.push(`PR size (${totalChanges} lines) exceeds ${MAX_LINES} line limit`);
               passed = false;
             }
-            
+
             // Check template
             const body = pr.body || '';
             const requiredSections = ["# Description", "Fixes #", "# How Has This Been Tested?", "# Checklist"];
             const missingSections = requiredSections.filter(section => !body.includes(section));
-            
+
             if (missingSections.length > 0) {
               errors.push(`Missing PR template sections: ${missingSections.join(', ')}`);
               passed = false;
             }
-            
+
             if (body.match(/Replace this with|Choose one:|Fixes #\[issue number\]/i)) {
               errors.push('PR template contains unmodified placeholders');
               passed = false;
             }
-            
+
             // Check linked issue
             const issueMatch = body.match(/(?:Fixes|Closes|Resolves)\s+#(\d+)/i);
             if (!issueMatch) {
               errors.push('No linked issue found');
               passed = false;
             }
-            
+
             // Store results
             core.setOutput('passed', passed);
             core.setOutput('errors', errors.join('; '));
             core.setOutput('totalChanges', totalChanges);
             core.setOutput('hasTemplate', missingSections.length === 0);
             core.setOutput('hasIssue', !!issueMatch);
-            
+
             if (!passed) {
               core.setFailed(errors.join('; '));
             }
@@ -109,7 +109,7 @@ jobs:
           script: |
             const passed = ${{ steps.validate.outputs.passed }};
             const errors = '${{ steps.validate.outputs.errors }}';
-            
+
             await github.rest.repos.createCommitStatus({
               owner: context.repo.owner,
               repo: context.repo.repo,
@@ -131,27 +131,27 @@ jobs:
             const hasTemplate = ${{ steps.validate.outputs.hasTemplate }};
             const hasIssue = ${{ steps.validate.outputs.hasIssue }};
             const errors = '${{ steps.validate.outputs.errors }}'.split('; ').filter(e => e);
-            
+
             let body = `### Manual Validation Results\n\n`;
             body += `**Status**: ${passed ? '✅ Passed' : '❌ Failed'}\n\n`;
             body += `| Check | Status | Details |\n`;
             body += `|-------|--------|----------|\n`;
             body += `| PR Size | ${totalChanges <= 1000 ? '✅' : '❌'} | ${totalChanges} lines ${totalChanges > 1000 ? '(exceeds 1000 limit)' : ''} |\n`;
             body += `| Template | ${hasTemplate ? '✅' : '❌'} | ${hasTemplate ? 'Complete' : 'Missing required sections'} |\n`;
             body += `| Linked Issue | ${hasIssue ? '✅' : '❌'} | ${hasIssue ? 'Found' : 'Missing Fixes/Closes #XXX'} |\n`;
-            
+
             if (errors.length > 0) {
               body += `\n**Errors:**\n`;
               errors.forEach(error => {
                 body += `- ❌ ${error}\n`;
               });
             }
-            
+
             body += `\n[View workflow run](https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId})`;
-            
+
             await github.rest.issues.createComment({
               owner: context.repo.owner,
               repo: context.repo.repo,
               issue_number: pr.number,
               body: body
-            });
+            });

langextract/visualization.py

@@ -28,7 +28,6 @@
 import html
 import itertools
 import json
-import os
 import pathlib
 import textwrap
 
@@ -37,10 +36,30 @@
 
 # Fallback if IPython is not present
 try:
-  from IPython.display import HTML  # type: ignore
-except Exception:
+  from IPython import get_ipython  # type: ignore[import-not-found]
+  from IPython.display import HTML  # type: ignore[import-not-found]
+except ImportError:
+
+  def get_ipython():  # type: ignore[no-redef]
+    return None
+
   HTML = None  # pytype: disable=annotation-type-mismatch
 
+
+def _is_jupyter() -> bool:
+  """Check if we're in a Jupyter/IPython environment that can display HTML."""
+  try:
+    if get_ipython is None:
+      return False
+    ip = get_ipython()
+    if ip is None:
+      return False
+    # Simple check: if we're in IPython and NOT in a plain terminal
+    return ip.__class__.__name__ != 'TerminalInteractiveShell'
+  except Exception:
+    return False
+
+
 _PALETTE: list[str] = [
     '#D2E3FC',  # Light Blue (Primary Container)
     '#C8E6C9',  # Light Green (Tertiary Container)
@@ -538,7 +557,7 @@ def visualize(
     animation_speed: float = 1.0,
     show_legend: bool = True,
     gif_optimized: bool = True,
-) -> 'HTML | str':
+) -> HTML | str:
   """Visualises extraction data as animated highlighted HTML.
 
   Args:
@@ -582,7 +601,9 @@ def visualize(
         ' animate.</p></div>'
     )
     full_html = _VISUALIZATION_CSS + empty_html
-    return HTML(full_html) if HTML is not None else full_html
+    if HTML is not None and _is_jupyter():
+      return HTML(full_html)
+    return full_html
 
   color_map = _assign_colors(valid_extractions)
 
@@ -603,4 +624,6 @@ def visualize(
         'class="lx-animated-wrapper lx-gif-optimized"',
     )
 
-  return HTML(full_html) if HTML is not None else full_html
+  if HTML is not None and _is_jupyter():
+    return HTML(full_html)
+  return full_html

pyproject.toml

@@ -63,6 +63,10 @@ test = [
     "pytest>=7.4.0",
     "tomli>=2.0.0"
 ]
+notebook = [
+    "ipython>=7.0.0",
+    "notebook>=6.0.0"
+]
 
 [tool.setuptools]
 packages = ["langextract"]